BladeDISC DeePMD code

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_tensor.py

+import os
+from typing import List, Optional, TYPE_CHECKING, Tuple
+
+import numpy as np
+from deepmd.common import make_default_mesh
+from deepmd.env import default_tf_session_config, tf
+from deepmd.infer.deep_eval import DeepEval
+from deepmd.utils.sess import run_sess
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+class DeepTensor(DeepEval):
+    """Evaluates a tensor model.
+
+    Parameters
+    ----------
+    model_file: str
+        The name of the frozen model file.
+    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 = {
+        # descriptor attrs
+        "t_ntypes": "descrpt_attr/ntypes:0",
+        "t_rcut": "descrpt_attr/rcut:0",
+        # model attrs
+        "t_tmap": "model_attr/tmap:0",
+        "t_sel_type": "model_attr/sel_type:0",
+        "t_ouput_dim": "model_attr/output_dim: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",
+        load_prefix: str = 'load',
+        default_tf_graph: bool = False
+    ) -> None:
+        """Constructor"""
+        DeepEval.__init__(
+            self,
+            model_file,
+            load_prefix=load_prefix,
+            default_tf_graph=default_tf_graph
+        )
+        # check model type
+        model_type = self.tensors["t_tensor"][2:-2]
+        assert self.model_type == model_type, \
+            f"expect {model_type} model but got {self.model_type}"
+
+        # now load tensors to object attributes
+        for attr_name, tensor_name in self.tensors.items():
+            self._get_tensor(tensor_name, attr_name)
+        
+        # load optional tensors if possible
+        optional_tensors = {
+            "t_global_tensor": f"o_global_{model_type}:0",
+            "t_force": "o_force:0",
+            "t_virial": "o_virial:0",
+            "t_atom_virial": "o_atom_virial:0"
+        }
+        try:
+            # first make sure these tensor all exists (but do not modify self attr)
+            for attr_name, tensor_name in optional_tensors.items():
+                self._get_tensor(tensor_name)
+            # then put those into self.attrs
+            for attr_name, tensor_name in optional_tensors.items():
+                self._get_tensor(tensor_name, attr_name)
+        except KeyError:
+            self._support_gfv = False
+        else:
+            self.tensors.update(optional_tensors)
+            self._support_gfv = True
+            
+        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.output_dim] \
+            = run_sess(self.sess, 
+                [self.t_ntypes, self.t_rcut, self.t_tmap, self.t_sel_type, self.t_ouput_dim]
+            )
+
+    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.ndarray,
+        cells: np.ndarray,
+        atom_types: List[int],
+        atomic: bool = True,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+        efield: Optional[np.ndarray] = None
+    ) -> np.ndarray:
+        """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
+            If True (default), return the atomic tensor
+            Otherwise return the global tensor
+        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 output_dim
+                else of size nframes x natoms x output_dim
+        """
+        # standarize the shape of inputs
+        atom_types = np.array(atom_types, dtype = int).reshape([-1])
+        natoms = atom_types.size
+        coords = np.reshape(np.array(coords), [-1, natoms * 3])
+        nframes = coords.shape[0]
+        if cells is None:
+            pbc = False
+            cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
+        else:
+            pbc = True
+            cells = np.array(cells).reshape([nframes, 9])
+
+        # 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
+        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])
+        feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
+        feed_dict_test[self.t_box  ] = np.reshape(cells , [-1])
+        if pbc:
+            feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
+        else:
+            feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
+
+        if atomic:
+            assert "global" not in self.model_type, \
+                f"cannot do atomic evaluation with model type {self.model_type}"
+            t_out = [self.t_tensor]
+        else:
+            assert self._support_gfv or "global" in self.model_type, \
+                f"do not support global tensor evaluation with old {self.model_type} model"
+            t_out = [self.t_global_tensor if self._support_gfv else self.t_tensor]
+        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.output_dim]), sel_imap)
+            tensor = np.reshape(tensor, [nframes, len(sel_at), self.output_dim])
+        else:
+            tensor = np.reshape(tensor, [nframes, self.output_dim])
+        
+        return tensor
+
+    def eval_full(
+        self,
+        coords: np.ndarray,
+        cells: np.ndarray,
+        atom_types: List[int],
+        atomic: bool = False,
+        fparam: Optional[np.array] = None,
+        aparam: Optional[np.array] = None,
+        efield: Optional[np.array] = None
+    ) -> Tuple[np.ndarray, ...]:
+        """Evaluate the model with interface similar to the energy model.
+        Will return global tensor, component-wise force and virial
+        and optionally atomic tensor and atomic virial.
+
+        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
+            Whether to calculate atomic tensor and virial
+        fparam
+            Not used in this model
+        aparam
+            Not used in this model
+        efield
+            Not used in this model
+
+        Returns
+        -------
+        tensor
+            The global tensor. 
+            shape: [nframes x nout]
+        force
+            The component-wise force (negative derivative) on each atom.
+            shape: [nframes x nout x natoms x 3]
+        virial
+            The component-wise virial of the tensor.
+            shape: [nframes x nout x 9]
+        atom_tensor
+            The atomic tensor. Only returned when atomic == True
+            shape: [nframes x natoms x nout]
+        atom_virial
+            The atomic virial. Only returned when atomic == True
+            shape: [nframes x nout x natoms x 9]
+        """
+        assert self._support_gfv, \
+            f"do not support eval_full with old tensor model"
+
+        # standarize the shape of inputs
+        atom_types = np.array(atom_types, dtype = int).reshape([-1])
+        natoms = atom_types.size
+        coords = np.reshape(np.array(coords), [-1, natoms * 3])
+        nframes = coords.shape[0]
+        if cells is None:
+            pbc = False
+            cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
+        else:
+            pbc = True
+            cells = np.array(cells).reshape([nframes, 9])
+        nout = self.output_dim
+
+        # 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
+        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])
+        feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
+        feed_dict_test[self.t_box  ] = np.reshape(cells , [-1])
+        if pbc:
+            feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
+        else:
+            feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
+        
+        t_out = [self.t_global_tensor, 
+                 self.t_force, 
+                 self.t_virial]
+        if atomic :
+            t_out += [self.t_tensor, 
+                      self.t_atom_virial]
+        
+        v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
+        gt     = v_out[0] # global tensor
+        force  = v_out[1]
+        virial = v_out[2]
+        if atomic:
+            at = v_out[3] # atom tensor
+            av = v_out[4] # atom virial
+
+        # please note here the shape are wrong!
+        force  = self.reverse_map(np.reshape(force, [nframes*nout, natoms ,3]), imap)
+        if atomic:
+            at = self.reverse_map(np.reshape(at, [nframes, len(sel_at), nout]), sel_imap)
+            av = self.reverse_map(np.reshape(av, [nframes*nout, natoms, 9]), imap)
+        
+        # make sure the shapes are correct here
+        gt     = np.reshape(gt, [nframes, nout])
+        force  = np.reshape(force, [nframes, nout, natoms, 3])
+        virial = np.reshape(virial, [nframes, nout, 9])
+        if atomic:
+            at = np.reshape(at, [nframes, len(sel_at), self.output_dim])
+            av = np.reshape(av, [nframes, nout, natoms, 9])
+            return gt, force, virial, at, av
+        else:
+            return gt, force, virial
\ No newline at end of file

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/deep_wfc.py

+from deepmd.infer.deep_tensor import DeepTensor
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+
+class DeepWFC(DeepTensor):
+    """Constructor.
+
+    Parameters
+    ----------
+    model_file : Path
+        The name of the frozen model file.
+    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
+
+    Warnings
+    --------
+    For developers: `DeepTensor` initializer must be called at the end after
+    `self.tensors` are modified because it uses the data in `self.tensors` dict.
+    Do not chanage the order!
+    """
+
+    def __init__(
+        self, model_file: "Path", load_prefix: str = "load", default_tf_graph: bool = False
+    ) -> None:
+
+        # use this in favor of dict update to move attribute from class to
+        # instance namespace
+        self.tensors = dict(
+            {
+                # output tensor
+                "t_tensor": "o_wfc:0",
+            },
+            **self.tensors
+        )
+        DeepTensor.__init__(
+            self,
+            model_file,
+            load_prefix=load_prefix,
+            default_tf_graph=default_tf_graph,
+        )
+
+    def get_dim_fparam(self) -> int:
+        """Unsupported in this model."""
+        raise NotImplementedError("This model type does not support this attribute")
+
+    def get_dim_aparam(self) -> int:
+        """Unsupported in this model."""
+        raise NotImplementedError("This model type does not support this attribute")

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/ewald_recp.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import ClassArg
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import global_cvt_2_ener_float
+from deepmd.env import op_module
+from deepmd.env import default_tf_session_config
+from deepmd.utils.sess import run_sess
+
+class EwaldRecp () :
+    """
+    Evaluate the reciprocal part of the Ewald sum
+    """
+    def __init__(self, 
+                 hh,
+                 beta):
+        """
+        Constructor 
+
+        Parameters
+        ----------
+        hh
+                Grid spacing of the reciprocal part of Ewald sum. Unit: A
+        beta
+                Splitting parameter of the Ewald sum. Unit: A^{-1}
+        """
+        self.hh = hh
+        self.beta = beta
+        with tf.Graph().as_default() as graph:
+            # place holders
+            self.t_nloc       = tf.placeholder(tf.int32, [1], name = "t_nloc")
+            self.t_coord      = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_coord')
+            self.t_charge     = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_charge')
+            self.t_box        = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_box')
+            # output            
+            self.t_energy, self.t_force, self.t_virial \
+                = op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box, 
+                                       ewald_h = self.hh,
+                                       ewald_beta = self.beta)
+        self.sess = tf.Session(graph=graph, config=default_tf_session_config)
+
+    def eval(self, 
+             coord : np.ndarray, 
+             charge : np.ndarray, 
+             box : np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray] :
+        """
+        Evaluate 
+        
+        Parameters
+        ----------
+        coord
+                The coordinates of atoms
+        charge
+                The atomic charge
+        box
+                The simulation region. PBC is assumed
+
+        Returns
+        -------
+        e
+                The energy 
+        f
+                The force 
+        v
+                The virial 
+        """
+        coord = np.array(coord)
+        charge = np.array(charge)
+        box = np.array(box)
+        nframes = charge.shape[0]
+        natoms = charge.shape[1]
+        coord = np.reshape(coord, [nframes * 3 * natoms])
+        charge = np.reshape(charge, [nframes * natoms])
+        box = np.reshape(box, [nframes * 9])
+
+        [energy, force, virial] \
+            = run_sess(self.sess, [self.t_energy, self.t_force, self.t_virial], 
+                            feed_dict = {
+                                self.t_coord:  coord,
+                                self.t_charge: charge,
+                                self.t_box:    box,
+                                self.t_nloc:   [natoms],
+                            })
+
+        return energy, force, virial
+             
+             

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/infer/model_devi.py

+import numpy as np
+from .deep_pot import DeepPot
+from ..utils.data import DeepmdData
+from ..utils.batch_size import AutoBatchSize
+from deepmd.common import expand_sys_str
+        
+
+def calc_model_devi_f(fs: np.ndarray):
+    '''
+    Parameters
+    ----------
+    fs : numpy.ndarray
+        size of `n_models x n_frames x n_atoms x 3`
+    '''
+    fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1)
+    max_devi_f = np.max(fs_devi, axis=-1)
+    min_devi_f = np.min(fs_devi, axis=-1)
+    avg_devi_f = np.mean(fs_devi, axis=-1)
+    return max_devi_f, min_devi_f, avg_devi_f
+
+def calc_model_devi_e(es: np.ndarray):
+    '''
+    Parameters
+    ----------
+    es : numpy.ndarray
+        size of `n_models x n_frames x n_atoms
+    '''
+    es_devi = np.std(es, axis=0)
+    max_devi_e = np.max(es_devi, axis=1)
+    min_devi_e = np.min(es_devi, axis=1)
+    avg_devi_e = np.mean(es_devi, axis=1)
+    return max_devi_e, min_devi_e, avg_devi_e
+
+def calc_model_devi_v(vs: np.ndarray):
+    '''
+    Parameters
+    ----------
+    vs : numpy.ndarray
+        size of `n_models x n_frames x 9`
+    '''
+    vs_devi = np.std(vs, axis=0)
+    max_devi_v = np.max(vs_devi, axis=-1)
+    min_devi_v = np.min(vs_devi, axis=-1)
+    avg_devi_v = np.linalg.norm(vs_devi, axis=-1) / 3
+    return max_devi_v, min_devi_v, avg_devi_v
+
+def write_model_devi_out(devi: np.ndarray, fname: str, header: str=""):
+    '''
+    Parameters
+    ----------
+    devi : numpy.ndarray
+        the first column is the steps index
+    fname : str
+        the file name to dump
+    header : str, default=""
+        the header to dump
+    '''
+    assert devi.shape[1] == 7
+    header = "%s\n%10s" % (header, "step")
+    for item in 'vf':
+        header += "%19s%19s%19s" % (f"max_devi_{item}", f"min_devi_{item}", f"avg_devi_{item}")
+    with open(fname, "ab") as fp:
+        np.savetxt(fp,
+                   devi,
+                   fmt=['%12d'] + ['%19.6e' for _ in range(6)],
+                   delimiter='',
+                   header=header)
+    return devi
+
+def _check_tmaps(tmaps, ref_tmap=None):
+    '''
+    Check whether type maps are identical
+    '''
+    assert isinstance(tmaps, list)
+    if ref_tmap is None:
+        ref_tmap = tmaps[0]
+    assert isinstance(ref_tmap, list)
+
+    flag = True
+    for tmap in tmaps:
+        if tmap != ref_tmap:
+            flag = False
+            break
+    return flag
+
+def calc_model_devi(coord,
+                    box,
+                    atype,
+                    models,
+                    fname=None,
+                    frequency=1, 
+                    ):
+    '''
+    Python interface to calculate model deviation
+
+    Parameters
+    -----------
+    coord : numpy.ndarray, `n_frames x n_atoms x 3`
+        Coordinates of system to calculate
+    box : numpy.ndarray or None, `n_frames x 3 x 3`
+        Box to specify periodic boundary condition. If None, no pbc will be used
+    atype : numpy.ndarray, `n_atoms x 1`
+        Atom types
+    models : list of DeepPot models
+        Models used to evaluate deviation
+    fname : str or None
+        File to dump results, default None
+    frequency : int
+        Steps between frames (if the system is given by molecular dynamics engine), default 1
+    
+    Returns
+    -------
+    model_devi : numpy.ndarray, `n_frames x 7`
+        Model deviation results. The first column is index of steps, the other 6 columns are
+        max_devi_v, min_devi_v, avg_devi_v, max_devi_f, min_devi_f, avg_devi_f.
+    
+    Examples
+    --------
+    >>> from deepmd.infer import calc_model_devi
+    >>> from deepmd.infer import DeepPot as DP
+    >>> import numpy as np
+    >>> coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
+    >>> cell = np.diag(10 * np.ones(3)).reshape([1, -1])
+    >>> atype = [1,0,1]
+    >>> graphs = [DP("graph.000.pb"), DP("graph.001.pb")]
+    >>> model_devi = calc_model_devi(coord, cell, atype, graphs)
+    '''
+    if box is not None:
+        nopbc = True
+    else:
+        nopbc = False
+
+    forces = []
+    virials = []
+    for dp in models:
+        ret = dp.eval(
+            coord,
+            box,
+            atype,
+        )
+        forces.append(ret[1])
+        virials.append(ret[2] / len(atype))
+    
+    forces = np.array(forces)
+    virials = np.array(virials)
+    
+    devi = [np.arange(coord.shape[0]) * frequency]
+    devi += list(calc_model_devi_v(virials))
+    devi += list(calc_model_devi_f(forces))
+    devi = np.vstack(devi).T
+    if fname:
+        write_model_devi_out(devi, fname)
+    return devi
+    
+def make_model_devi(
+    *,
+    models: list,
+    system: str,
+    set_prefix: str,
+    output: str,
+    frequency: int,
+    **kwargs
+):
+    '''
+    Make model deviation calculation
+
+    Parameters
+    ----------
+    models: list
+        A list of paths of models to use for making model deviation
+    system: str
+        The path of system to make model deviation calculation
+    set_prefix: str
+        The set prefix of the system
+    output: str
+        The output file for model deviation results
+    frequency: int
+        The number of steps that elapse between writing coordinates 
+        in a trajectory by a MD engine (such as Gromacs / Lammps).
+        This paramter is used to determine the index in the output file.
+    '''
+    auto_batch_size = AutoBatchSize()
+    # init models
+    dp_models = [DeepPot(model, auto_batch_size=auto_batch_size) for model in models]
+
+    # check type maps
+    tmaps = [dp.get_type_map() for dp in dp_models]
+    if _check_tmaps(tmaps):
+        tmap = tmaps[0]
+    else:
+        raise RuntimeError("The models does not have the same type map.")
+
+    all_sys = expand_sys_str(system)
+    if len(all_sys) == 0:
+        raise RuntimeError("Did not find valid system")
+    devis_coll = []
+    for system in all_sys:
+        # create data-system
+        dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
+
+        data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
+        nframes_tot = 0
+        devis = []
+        for data in data_sets:
+            coord = data["coord"]
+            box = data["box"]
+            atype = data["type"][0] 
+            if not dp_data.pbc:
+                box = None
+            devi = calc_model_devi(coord, box, atype, dp_models)
+            nframes_tot += coord.shape[0]
+            devis.append(devi)
+        devis = np.vstack(devis)
+        devis[:, 0] = np.arange(nframes_tot) * frequency
+        write_model_devi_out(devis, output, header=system)
+        devis_coll.append(devis)
+    return devis_coll

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loggers/__init__.py

+"""Module taking care of logging duties."""
+
+from .loggers import set_log_handles
+
+__all__ = ["set_log_handles"]

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loggers/loggers.py

+"""Logger initialization for package."""
+
+import logging
+import os
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from pathlib import Path
+
+    from mpi4py import MPI
+
+    _MPI_APPEND_MODE = MPI.MODE_CREATE | MPI.MODE_APPEND
+
+logging.getLogger(__name__)
+
+__all__ = ["set_log_handles"]
+
+# logger formater
+FFORMATTER = logging.Formatter(
+    "[%(asctime)s] %(app_name)s %(levelname)-7s %(name)-45s %(message)s"
+)
+CFORMATTER = logging.Formatter(
+#    "%(app_name)s %(levelname)-7s |-> %(name)-45s %(message)s"
+    "%(app_name)s %(levelname)-7s %(message)s"
+)
+FFORMATTER_MPI = logging.Formatter(
+    "[%(asctime)s] %(app_name)s rank:%(rank)-2s %(levelname)-7s %(name)-45s %(message)s"
+)
+CFORMATTER_MPI = logging.Formatter(
+#    "%(app_name)s rank:%(rank)-2s %(levelname)-7s |-> %(name)-45s %(message)s"
+    "%(app_name)s rank:%(rank)-2s %(levelname)-7s %(message)s"
+)
+
+
+class _AppFilter(logging.Filter):
+    """Add field `app_name` to log messages."""
+
+    def filter(self, record):
+        record.app_name = "DEEPMD"
+        return True
+
+
+class _MPIRankFilter(logging.Filter):
+    """Add MPI rank number to log messages, adds field `rank`."""
+
+    def __init__(self, rank: int) -> None:
+        super().__init__(name="MPI_rank_id")
+        self.mpi_rank = str(rank)
+
+    def filter(self, record):
+        record.rank = self.mpi_rank
+        return True
+
+
+class _MPIMasterFilter(logging.Filter):
+    """Filter that lets through only messages emited from rank==0."""
+
+    def __init__(self, rank: int) -> None:
+        super().__init__(name="MPI_master_log")
+        self.mpi_rank = rank
+
+    def filter(self, record):
+        if self.mpi_rank == 0:
+            return True
+        else:
+            return False
+
+
+class _MPIFileStream:
+    """Wrap MPI.File` so it has the same API as python file streams.
+
+    Parameters
+    ----------
+    filename : Path
+        disk location of the file stream
+    MPI : MPI
+        MPI communicator object
+    mode : str, optional
+        file write mode, by default _MPI_APPEND_MODE
+    """
+
+    def __init__(
+        self, filename: "Path", MPI: "MPI", mode: str = "_MPI_APPEND_MODE"
+    ) -> None:
+        self.stream = MPI.File.Open(MPI.COMM_WORLD, filename, mode)
+        self.stream.Set_atomicity(True)
+        self.name = "MPIfilestream"
+
+    def write(self, msg: str):
+        """Write to MPI shared file stream.
+
+        Parameters
+        ----------
+        msg : str
+            message to write
+        """
+        b = bytearray()
+        b.extend(map(ord, msg))
+        self.stream.Write_shared(b)
+
+    def close(self):
+        """Synchronize and close MPI file stream."""
+        self.stream.Sync()
+        self.stream.Close()
+
+
+class _MPIHandler(logging.FileHandler):
+    """Emulate `logging.FileHandler` with MPI shared File that all ranks can write to.
+
+    Parameters
+    ----------
+    filename : Path
+        file path
+    MPI : MPI
+        MPI communicator object
+    mode : str, optional
+        file access mode, by default "_MPI_APPEND_MODE"
+    """
+
+    def __init__(
+        self,
+        filename: "Path",
+        MPI: "MPI",
+        mode: str = "_MPI_APPEND_MODE",
+    ) -> None:
+        self.MPI = MPI
+        super().__init__(filename, mode=mode, encoding=None, delay=False)
+
+    def _open(self):
+        return _MPIFileStream(self.baseFilename, self.MPI, self.mode)
+
+    def setStream(self, stream):
+        """Stream canot be reasigned in MPI mode."""
+        raise NotImplementedError("Unable to do for MPI file handler!")
+
+
+def set_log_handles(
+    level: int,
+    log_path: Optional["Path"] = None,
+    mpi_log: Optional[str] = None
+):
+    """Set desired level for package loggers and add file handlers.
+
+    Parameters
+    ----------
+    level: int
+        logging level
+    log_path: Optional[str]
+        path to log file, if None logs will be send only to console. If the parent
+        directory does not exist it will be automatically created, by default None
+    mpi_log : Optional[str], optional
+        mpi log type. Has three options. `master` will output logs to file and console
+        only from rank==0. `collect` will write messages from all ranks to one file
+        opened under rank==0 and to console. `workers` will open one log file for each
+        worker designated by its rank, console behaviour is the same as for `collect`.
+        If this argument is specified, package 'mpi4py' must be already installed.
+        by default None
+
+    Raises
+    ------
+    RuntimeError
+        If the argument `mpi_log` is specified, package `mpi4py` is not installed.
+
+    References
+    ----------
+    https://groups.google.com/g/mpi4py/c/SaNzc8bdj6U
+    https://stackoverflow.com/questions/35869137/avoid-tensorflow-print-on-standard-error
+    https://stackoverflow.com/questions/56085015/suppress-openmp-debug-messages-when-running-tensorflow-on-cpu
+
+    Notes
+    -----
+    Logging levels:
+
+    +---------+--------------+----------------+----------------+----------------+
+    |         | our notation | python logging | tensorflow cpp | OpenMP         |
+    +=========+==============+================+================+================+
+    | debug   | 10           | 10             | 0              | 1/on/true/yes  |
+    +---------+--------------+----------------+----------------+----------------+
+    | info    | 20           | 20             | 1              | 0/off/false/no |
+    +---------+--------------+----------------+----------------+----------------+
+    | warning | 30           | 30             | 2              | 0/off/false/no |
+    +---------+--------------+----------------+----------------+----------------+
+    | error   | 40           | 40             | 3              | 0/off/false/no |
+    +---------+--------------+----------------+----------------+----------------+
+
+    """
+    # silence logging for OpenMP when running on CPU if level is any other than debug
+    if level <= 10:
+        os.environ["KMP_WARNINGS"] = "FALSE"
+
+    # set TF cpp internal logging level
+    os.environ['TF_CPP_MIN_LOG_LEVEL'] = str(int((level / 10) - 1))
+
+    # get root logger
+    root_log = logging.getLogger()
+
+    # remove all old handlers
+    root_log.setLevel(level)
+    for hdlr in root_log.handlers[:]:
+        root_log.removeHandler(hdlr)
+
+    # check if arguments are present
+    MPI = None
+    if mpi_log:
+        try:
+            from mpi4py import MPI
+        except ImportError as e:
+            raise RuntimeError("You cannot specify 'mpi_log' when mpi4py not installed") from e
+
+    # * add console handler ************************************************************
+    ch = logging.StreamHandler()
+    if MPI:
+        rank = MPI.COMM_WORLD.Get_rank()
+        if mpi_log == "master":
+            ch.setFormatter(CFORMATTER)
+            ch.addFilter(_MPIMasterFilter(rank))
+        else:
+            ch.setFormatter(CFORMATTER_MPI)
+            ch.addFilter(_MPIRankFilter(rank))
+    else:
+        ch.setFormatter(CFORMATTER)
+
+    ch.setLevel(level)
+    ch.addFilter(_AppFilter())
+    root_log.addHandler(ch)
+
+    # * add file handler ***************************************************************
+    if log_path:
+
+        # create directory
+        log_path.parent.mkdir(exist_ok=True, parents=True)
+
+        fh = None
+
+        if mpi_log == "master":
+            rank = MPI.COMM_WORLD.Get_rank()
+            if rank == 0:
+                fh = logging.FileHandler(log_path, mode="w")
+                fh.addFilter(_MPIMasterFilter(rank))
+                fh.setFormatter(FFORMATTER)
+        elif mpi_log == "collect":
+            rank = MPI.COMM_WORLD.Get_rank()
+            fh = _MPIHandler(log_path, MPI, mode=MPI.MODE_WRONLY | MPI.MODE_CREATE)
+            fh.addFilter(_MPIRankFilter(rank))
+            fh.setFormatter(FFORMATTER_MPI)
+        elif mpi_log == "workers":
+            rank = MPI.COMM_WORLD.Get_rank()
+            # if file has suffix than inser rank number before suffix
+            # e.g deepmd.log -> deepmd_<rank>.log
+            # if no suffix is present, insert rank as suffix
+            # e.g. deepmdlog -> deepmdlog.<rank>
+            if log_path.suffix:
+                worker_log = (log_path.parent / f"{log_path.stem}_{rank}").with_suffix(
+                    log_path.suffix
+                )
+            else:
+                worker_log = log_path.with_suffix(f".{rank}")
+
+            fh = logging.FileHandler(worker_log, mode="w")
+            fh.setFormatter(FFORMATTER)
+        else:
+            fh = logging.FileHandler(log_path, mode="w")
+            fh.setFormatter(FFORMATTER)
+
+        if fh:
+            fh.setLevel(level)
+            fh.addFilter(_AppFilter())
+            root_log.addHandler(fh)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loss/__init__.py

+from .ener import EnerStdLoss
+from .ener import EnerDipoleLoss
+from .tensor import TensorLoss
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loss/ener.py

+import numpy as np
+from deepmd.env import tf
+from deepmd.common import ClassArg, add_data_requirement
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import global_cvt_2_ener_float
+from deepmd.utils.sess import run_sess
+from .loss import Loss
+
+
+class EnerStdLoss (Loss) :
+    """
+    Standard loss function for DP models
+
+    Parameters
+    ----------
+    enable_atom_ener_coeff : bool
+        if true, the energy will be computed as \sum_i c_i E_i
+    """
+    def __init__ (self, 
+                  starter_learning_rate : float, 
+                  start_pref_e : float = 0.02,
+                  limit_pref_e : float = 1.00,
+                  start_pref_f : float = 1000,
+                  limit_pref_f : float = 1.00,
+                  start_pref_v : float = 0.0,
+                  limit_pref_v : float = 0.0,
+                  start_pref_ae : float = 0.0,
+                  limit_pref_ae : float = 0.0,
+                  start_pref_pf : float = 0.0,
+                  limit_pref_pf : float = 0.0,
+                  relative_f : float = None,
+                  enable_atom_ener_coeff: bool=False,
+    ) -> None:
+        self.starter_learning_rate = starter_learning_rate
+        self.start_pref_e = start_pref_e
+        self.limit_pref_e = limit_pref_e
+        self.start_pref_f = start_pref_f
+        self.limit_pref_f = limit_pref_f
+        self.start_pref_v = start_pref_v
+        self.limit_pref_v = limit_pref_v
+        self.start_pref_ae = start_pref_ae
+        self.limit_pref_ae = limit_pref_ae
+        self.start_pref_pf = start_pref_pf
+        self.limit_pref_pf = limit_pref_pf
+        self.relative_f = relative_f
+        self.enable_atom_ener_coeff = enable_atom_ener_coeff
+        self.has_e = (self.start_pref_e != 0.0 or self.limit_pref_e != 0.0)
+        self.has_f = (self.start_pref_f != 0.0 or self.limit_pref_f != 0.0)
+        self.has_v = (self.start_pref_v != 0.0 or self.limit_pref_v != 0.0)
+        self.has_ae = (self.start_pref_ae != 0.0 or self.limit_pref_ae != 0.0)
+        self.has_pf = (self.start_pref_pf != 0.0 or self.limit_pref_pf != 0.0)
+        # data required
+        add_data_requirement('energy', 1, atomic=False, must=False, high_prec=True)
+        add_data_requirement('force',  3, atomic=True,  must=False, high_prec=False)
+        add_data_requirement('virial', 9, atomic=False, must=False, high_prec=False)
+        add_data_requirement('atom_ener', 1, atomic=True, must=False, high_prec=False)
+        add_data_requirement('atom_pref', 1, atomic=True, must=False, high_prec=False, repeat=3)
+        if self.enable_atom_ener_coeff:
+            add_data_requirement('atom_ener_coeff', 1, atomic=True, must=False, high_prec=False, default=1.)
+
+    def build (self, 
+               learning_rate,
+               natoms,
+               model_dict,
+               label_dict,
+               suffix):        
+        energy = model_dict['energy']
+        force = model_dict['force']
+        virial = model_dict['virial']
+        atom_ener = model_dict['atom_ener']
+        energy_hat = label_dict['energy']
+        force_hat = label_dict['force']
+        virial_hat = label_dict['virial']
+        atom_ener_hat = label_dict['atom_ener']
+        atom_pref = label_dict['atom_pref']
+        find_energy = label_dict['find_energy']
+        find_force = label_dict['find_force']
+        find_virial = label_dict['find_virial']
+        find_atom_ener = label_dict['find_atom_ener']                
+        find_atom_pref = label_dict['find_atom_pref']                
+
+        if self.enable_atom_ener_coeff:
+            # when ener_coeff (\nu) is defined, the energy is defined as 
+            # E = \sum_i \nu_i E_i
+            # instead of the sum of atomic energies.
+            #
+            # A case is that we want to train reaction energy
+            # A + B -> C + D
+            # E = - E(A) - E(B) + E(C) + E(D)
+            # A, B, C, D could be put far away from each other
+            atom_ener_coeff = label_dict['atom_ener_coeff']
+            atom_ener_coeff = tf.reshape(atom_ener_coeff, tf.shape(atom_ener))
+            energy = tf.reduce_sum(atom_ener_coeff * atom_ener, 1)
+        l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
+
+        force_reshape = tf.reshape (force, [-1])
+        force_hat_reshape = tf.reshape (force_hat, [-1])
+        atom_pref_reshape = tf.reshape (atom_pref, [-1])
+        diff_f = force_hat_reshape - force_reshape
+        if self.relative_f is not None:            
+            force_hat_3 = tf.reshape(force_hat, [-1, 3])
+            norm_f = tf.reshape(tf.norm(force_hat_3, axis = 1), [-1, 1]) + self.relative_f
+            diff_f_3 = tf.reshape(diff_f, [-1, 3])
+            diff_f_3 = diff_f_3 / norm_f
+            diff_f = tf.reshape(diff_f_3, [-1])
+        l2_force_loss = tf.reduce_mean(tf.square(diff_f), name = "l2_force_" + suffix)
+        l2_pref_force_loss = tf.reduce_mean(tf.multiply(tf.square(diff_f), atom_pref_reshape), name = "l2_pref_force_" + suffix)
+
+        virial_reshape = tf.reshape (virial, [-1])
+        virial_hat_reshape = tf.reshape (virial_hat, [-1])
+        l2_virial_loss = tf.reduce_mean (tf.square(virial_hat_reshape - virial_reshape), name = "l2_virial_" + suffix)
+
+        atom_ener_reshape = tf.reshape (atom_ener, [-1])
+        atom_ener_hat_reshape = tf.reshape (atom_ener_hat, [-1])
+        l2_atom_ener_loss = tf.reduce_mean (tf.square(atom_ener_hat_reshape - atom_ener_reshape), name = "l2_atom_ener_" + suffix)
+
+        atom_norm  = 1./ global_cvt_2_tf_float(natoms[0]) 
+        atom_norm_ener  = 1./ global_cvt_2_ener_float(natoms[0]) 
+        pref_e = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
+        pref_f = global_cvt_2_tf_float(find_force * (self.limit_pref_f + (self.start_pref_f - self.limit_pref_f) * learning_rate / self.starter_learning_rate) )
+        pref_v = global_cvt_2_tf_float(find_virial * (self.limit_pref_v + (self.start_pref_v - self.limit_pref_v) * learning_rate / self.starter_learning_rate) )
+        pref_ae= global_cvt_2_tf_float(find_atom_ener * (self.limit_pref_ae+ (self.start_pref_ae-self.limit_pref_ae) * learning_rate / self.starter_learning_rate) )
+        pref_pf= global_cvt_2_tf_float(find_atom_pref * (self.limit_pref_pf+ (self.start_pref_pf-self.limit_pref_pf) * learning_rate / self.starter_learning_rate) )
+
+        l2_loss = 0
+        more_loss = {}
+        if self.has_e :
+            l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
+        more_loss['l2_ener_loss'] = l2_ener_loss
+        if self.has_f :
+            l2_loss += global_cvt_2_ener_float(pref_f * l2_force_loss)
+        more_loss['l2_force_loss'] = l2_force_loss
+        if self.has_v :
+            l2_loss += global_cvt_2_ener_float(atom_norm * (pref_v * l2_virial_loss))
+        more_loss['l2_virial_loss'] = l2_virial_loss
+        if self.has_ae :
+            l2_loss += global_cvt_2_ener_float(pref_ae * l2_atom_ener_loss)
+        more_loss['l2_atom_ener_loss'] = l2_atom_ener_loss
+        if self.has_pf :
+            l2_loss += global_cvt_2_ener_float(pref_pf * l2_pref_force_loss)
+        more_loss['l2_pref_force_loss'] = l2_pref_force_loss
+
+        # only used when tensorboard was set as true
+        self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
+        self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', global_cvt_2_tf_float(tf.sqrt(l2_ener_loss)) / global_cvt_2_tf_float(natoms[0]))
+        self.l2_loss_force_summary = tf.summary.scalar('l2_force_loss', tf.sqrt(l2_force_loss))
+        self.l2_loss_virial_summary = tf.summary.scalar('l2_virial_loss', tf.sqrt(l2_virial_loss) / global_cvt_2_tf_float(natoms[0]))
+
+        self.l2_l = l2_loss
+        self.l2_more = more_loss
+        return l2_loss, more_loss
+
+    def eval(self, sess, feed_dict, natoms):
+        placeholder = self.l2_l
+        run_data = [
+            self.l2_l,
+            self.l2_more['l2_ener_loss'] if self.has_e else placeholder,
+            self.l2_more['l2_force_loss'] if self.has_f else placeholder,
+            self.l2_more['l2_virial_loss'] if self.has_v else placeholder,
+            self.l2_more['l2_atom_ener_loss'] if self.has_ae else placeholder,
+            self.l2_more['l2_pref_force_loss'] if self.has_pf else placeholder,
+        ]
+        error, error_e, error_f, error_v, error_ae, error_pf = run_sess(sess, run_data, feed_dict=feed_dict)
+        results = {"natoms": natoms[0], "rmse": np.sqrt(error)}
+        if self.has_e:
+            results["rmse_e"] = np.sqrt(error_e) / natoms[0]
+        if self.has_ae:
+            results["rmse_ae"] = np.sqrt(error_ae)
+        if self.has_f:
+            results["rmse_f"] = np.sqrt(error_f)
+        if self.has_v:
+            results["rmse_v"] = np.sqrt(error_v) / natoms[0]
+        if self.has_pf:
+            results["rmse_pf"] = np.sqrt(error_pf)
+        return results
+
+    def print_header(self):  # depreciated
+        prop_fmt = '   %11s %11s'
+        print_str = ''
+        print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
+        if self.has_e :
+            print_str += prop_fmt % ('rmse_e_tst', 'rmse_e_trn')
+        if self.has_ae :
+            print_str += prop_fmt % ('rmse_ae_tst', 'rmse_ae_trn')
+        if self.has_f :
+            print_str += prop_fmt % ('rmse_f_tst', 'rmse_f_trn')
+        if self.has_v :
+            print_str += prop_fmt % ('rmse_v_tst', 'rmse_v_trn')
+        if self.has_pf :
+            print_str += prop_fmt % ('rmse_pf_tst', 'rmse_pf_trn')
+        return print_str
+
+    def print_on_training(self, 
+                          tb_writer,
+                          cur_batch,
+                          sess, 
+                          natoms,
+                          feed_dict_test,
+                          feed_dict_batch):  # depreciated
+
+        run_data = [
+            self.l2_l,
+            self.l2_more['l2_ener_loss'],
+            self.l2_more['l2_force_loss'],
+            self.l2_more['l2_virial_loss'],
+            self.l2_more['l2_atom_ener_loss'],
+            self.l2_more['l2_pref_force_loss']
+        ]
+
+        # first train data
+        train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
+        error_train, error_e_train, error_f_train, error_v_train, error_ae_train, error_pf_train = train_out
+
+        # than test data, if tensorboard log writter is present, commpute summary
+        # and write tensorboard logs
+        if tb_writer:
+            summary_merged_op = tf.summary.merge([self.l2_loss_summary, self.l2_loss_ener_summary, self.l2_loss_force_summary, self.l2_loss_virial_summary])
+            run_data.insert(0, summary_merged_op)
+
+        test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
+
+        if tb_writer:
+            summary = test_out.pop(0)
+            tb_writer.add_summary(summary, cur_batch)
+
+        error_test, error_e_test, error_f_test, error_v_test, error_ae_test, error_pf_test = test_out
+
+        
+        print_str = ""
+        prop_fmt = "   %11.2e %11.2e"
+        print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
+        if self.has_e :
+            print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
+        if self.has_ae :
+            print_str += prop_fmt % (np.sqrt(error_ae_test), np.sqrt(error_ae_train))
+        if self.has_f :
+            print_str += prop_fmt % (np.sqrt(error_f_test), np.sqrt(error_f_train))
+        if self.has_v :
+            print_str += prop_fmt % (np.sqrt(error_v_test) / natoms[0], np.sqrt(error_v_train) / natoms[0])
+        if self.has_pf:
+            print_str += prop_fmt % (np.sqrt(error_pf_test), np.sqrt(error_pf_train))
+
+        return print_str      
+
+
+class EnerDipoleLoss (Loss) :
+    def __init__ (self, 
+                  starter_learning_rate : float,
+                  start_pref_e : float = 0.1,
+                  limit_pref_e : float = 1.0,
+                  start_pref_ed : float = 1.0,
+                  limit_pref_ed : float = 1.0
+    ) -> None :
+        self.starter_learning_rate = kwarg['starter_learning_rate']
+        args = ClassArg()\
+            .add('start_pref_e',        float,  must = True, default = 0.1) \
+            .add('limit_pref_e',        float,  must = True, default = 1.00)\
+            .add('start_pref_ed',       float,  must = True, default = 1.00)\
+            .add('limit_pref_ed',       float,  must = True, default = 1.00)
+        class_data = args.parse(jdata)
+        self.start_pref_e = class_data['start_pref_e']
+        self.limit_pref_e = class_data['limit_pref_e']
+        self.start_pref_ed = class_data['start_pref_ed']
+        self.limit_pref_ed = class_data['limit_pref_ed']
+        # data required
+        add_data_requirement('energy', 1, atomic=False, must=True, high_prec=True)
+        add_data_requirement('energy_dipole', 3, atomic=False, must=True, high_prec=False)
+
+    def build (self, 
+               learning_rate,
+               natoms,
+               model_dict,
+               label_dict,
+               suffix):        
+        coord = model_dict['coord']
+        energy = model_dict['energy']
+        atom_ener = model_dict['atom_ener']
+        nframes = tf.shape(atom_ener)[0]
+        natoms = tf.shape(atom_ener)[1]
+        # build energy dipole
+        atom_ener0 = atom_ener - tf.reshape(tf.tile(tf.reshape(energy/global_cvt_2_ener_float(natoms), [-1, 1]), [1, natoms]), [nframes, natoms])
+        coord = tf.reshape(coord, [nframes, natoms, 3])
+        atom_ener0 = tf.reshape(atom_ener0, [nframes, 1, natoms])
+        ener_dipole = tf.matmul(atom_ener0, coord)
+        ener_dipole = tf.reshape(ener_dipole, [nframes, 3])
+        
+        energy_hat = label_dict['energy']
+        ener_dipole_hat = label_dict['energy_dipole']
+        find_energy = label_dict['find_energy']
+        find_ener_dipole = label_dict['find_energy_dipole']                
+
+        l2_ener_loss = tf.reduce_mean( tf.square(energy - energy_hat), name='l2_'+suffix)
+
+        ener_dipole_reshape = tf.reshape(ener_dipole, [-1])
+        ener_dipole_hat_reshape = tf.reshape(ener_dipole_hat, [-1])
+        l2_ener_dipole_loss = tf.reduce_mean( tf.square(ener_dipole_reshape - ener_dipole_hat_reshape), name='l2_'+suffix)
+
+        # atom_norm_ener  = 1./ global_cvt_2_ener_float(natoms[0]) 
+        atom_norm_ener  = 1./ global_cvt_2_ener_float(natoms) 
+        pref_e  = global_cvt_2_ener_float(find_energy * (self.limit_pref_e + (self.start_pref_e - self.limit_pref_e) * learning_rate / self.starter_learning_rate) )
+        pref_ed = global_cvt_2_tf_float(find_ener_dipole * (self.limit_pref_ed + (self.start_pref_ed - self.limit_pref_ed) * learning_rate / self.starter_learning_rate) )
+
+        l2_loss = 0
+        more_loss = {}
+        l2_loss += atom_norm_ener * (pref_e * l2_ener_loss)
+        l2_loss += global_cvt_2_ener_float(pref_ed * l2_ener_dipole_loss)
+        more_loss['l2_ener_loss'] = l2_ener_loss
+        more_loss['l2_ener_dipole_loss'] = l2_ener_dipole_loss
+
+        self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
+        self.l2_loss_ener_summary = tf.summary.scalar('l2_ener_loss', tf.sqrt(l2_ener_loss) / global_cvt_2_tf_float(natoms[0]))
+        self.l2_ener_dipole_loss_summary = tf.summary.scalar('l2_ener_dipole_loss', tf.sqrt(l2_ener_dipole_loss))
+
+        self.l2_l = l2_loss
+        self.l2_more = more_loss
+        return l2_loss, more_loss
+
+    def eval(self, sess, feed_dict, natoms):
+        run_data = [
+            self.l2_l,
+            self.l2_more['l2_ener_loss'],
+            self.l2_more['l2_ener_dipole_loss']
+        ]
+        error, error_e, error_ed = run_sess(sess, run_data, feed_dict=feed_dict)
+        results = {
+            'natoms': natoms[0],
+            'rmse': np.sqrt(error),
+            'rmse_e': np.sqrt(error_e) / natoms[0],
+            'rmse_ed': np.sqrt(error_ed)
+        }
+        return results
+
+    @staticmethod
+    def print_header() :  # depreciated
+        prop_fmt = '   %9s %9s'
+        print_str = ''
+        print_str += prop_fmt % ('l2_tst', 'l2_trn')
+        print_str += prop_fmt % ('l2_e_tst', 'l2_e_trn')
+        print_str += prop_fmt % ('l2_ed_tst', 'l2_ed_trn')
+        return print_str 
+
+    def print_on_training(self,
+                          tb_writer,
+                          cur_batch, 
+                          sess, 
+                          natoms,
+                          feed_dict_test,
+                          feed_dict_batch):  # depreciated
+
+        run_data = [
+            self.l2_l,
+            self.l2_more['l2_ener_loss'],
+            self.l2_more['l2_ener_dipole_loss']
+        ]
+
+        # first train data
+        train_out = run_sess(sess, run_data, feed_dict=feed_dict_batch)
+        error_train, error_e_train, error_ed_train = train_out
+
+        # than test data, if tensorboard log writter is present, commpute summary
+        # and write tensorboard logs
+        if tb_writer:
+            summary_merged_op = tf.summary.merge([
+                self.l2_loss_summary,
+                self.l2_loss_ener_summary,
+                self.l2_ener_dipole_loss_summary
+            ])
+            run_data.insert(0, summary_merged_op)
+
+        test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
+
+        if tb_writer:
+            summary = test_out.pop(0)
+            tb_writer.add_summary(summary, cur_batch)
+
+        error_test, error_e_test, error_ed_test = test_out
+
+        print_str = ""
+        prop_fmt = "   %9.2e %9.2e"
+        print_str += prop_fmt % (np.sqrt(error_test), np.sqrt(error_train))
+        print_str += prop_fmt % (np.sqrt(error_e_test) / natoms[0], np.sqrt(error_e_train) / natoms[0])
+        print_str += prop_fmt % (np.sqrt(error_ed_test), np.sqrt(error_ed_train))
+        return print_str   
+
+
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loss/loss.py

+from abc import ABCMeta, abstractmethod
+from typing import Tuple, Dict
+from deepmd.env import tf
+
+
+class Loss(metaclass=ABCMeta):
+    """The abstract class for the loss function."""
+    @abstractmethod
+    def build(self, 
+            learning_rate: tf.Tensor,
+            natoms: tf.Tensor,
+            model_dict: Dict[str, tf.Tensor],
+            label_dict: Dict[str, tf.Tensor],
+            suffix: str) -> Tuple[tf.Tensor, Dict[str, tf.Tensor]]:
+        """Build the loss function graph.
+        
+        Parameters
+        ----------
+        learning_rate : tf.Tensor
+            learning rate
+        natoms : tf.Tensor
+            number of atoms
+        model_dict : dict[str, tf.Tensor]
+            A dictionary that maps model keys to tensors
+        label_dict : dict[str, tf.Tensor]
+            A dictionary that maps label keys to tensors
+        suffix : str
+            suffix
+
+        Returns
+        -------
+        tf.Tensor
+            the total squared loss
+        dict[str, tf.Tensor]
+            A dictionary that maps loss keys to more loss tensors
+        """
+
+    @abstractmethod
+    def eval(self,
+             sess: tf.Session,
+             feed_dict: Dict[tf.placeholder, tf.Tensor],
+             natoms: tf.Tensor) -> dict:
+        """Eval the loss function.
+
+        Parameters
+        ----------
+        sess : tf.Session
+            TensorFlow session
+        feed_dict : dict[tf.placeholder, tf.Tensor]
+            A dictionary that maps graph elements to values
+        natoms : tf.Tensor
+            number of atoms
+
+        Returns
+        -------
+        dict
+            A dictionary that maps keys to values. It
+            should contain key `natoms`
+        """

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/loss/tensor.py

+import numpy as np
+from deepmd.env import tf
+from deepmd.common import ClassArg, add_data_requirement
+
+from deepmd.env import global_cvt_2_tf_float
+from deepmd.env import global_cvt_2_ener_float
+from deepmd.utils.sess import run_sess
+from .loss import Loss
+
+
+class TensorLoss(Loss) :
+    """
+    Loss function for tensorial properties.
+    """
+    def __init__ (self, jdata, **kwarg) :
+        model = kwarg.get('model', None)
+        if model is not None:
+            self.type_sel = model.get_sel_type()
+        else:
+            self.type_sel = None
+        self.tensor_name = kwarg['tensor_name']
+        self.tensor_size = kwarg['tensor_size']
+        self.label_name = kwarg['label_name']
+        if jdata is not None:
+            self.scale = jdata.get('scale', 1.0)
+        else:
+            self.scale = 1.0
+
+        # YHT: added for global / local dipole combination
+        assert jdata is not None, "Please provide loss parameters!"
+        # YWolfeee: modify, use pref / pref_atomic, instead of pref_weight / pref_atomic_weight
+        self.local_weight = jdata.get('pref_atomic', None)
+        self.global_weight = jdata.get('pref', None)
+
+        assert (self.local_weight is not None and self.global_weight is not None), "Both `pref` and `pref_atomic` should be provided."
+        assert self.local_weight >= 0.0 and self.global_weight >= 0.0, "Can not assign negative weight to `pref` and `pref_atomic`"
+        assert (self.local_weight >0.0) or (self.global_weight>0.0), AssertionError('Can not assian zero weight both to `pref` and `pref_atomic`')
+
+        # data required
+        add_data_requirement("atomic_" + self.label_name, 
+                             self.tensor_size, 
+                             atomic=True,  
+                             must=False, 
+                             high_prec=False, 
+                             type_sel = self.type_sel)
+        add_data_requirement(self.label_name, 
+                             self.tensor_size, 
+                             atomic=False,  
+                             must=False, 
+                             high_prec=False, 
+                             type_sel = self.type_sel)
+
+    def build (self, 
+               learning_rate,
+               natoms,
+               model_dict,
+               label_dict,
+               suffix):        
+        polar_hat = label_dict[self.label_name]
+        atomic_polar_hat = label_dict["atomic_" + self.label_name]
+        polar = tf.reshape(model_dict[self.tensor_name], [-1])
+
+        find_global = label_dict['find_' + self.label_name]
+        find_atomic = label_dict['find_atomic_' + self.label_name]
+        
+        
+
+        # YHT: added for global / local dipole combination
+        l2_loss = global_cvt_2_tf_float(0.0)
+        more_loss = {
+            "local_loss":global_cvt_2_tf_float(0.0),
+            "global_loss":global_cvt_2_tf_float(0.0)
+        }
+
+        
+        if self.local_weight > 0.0:
+            local_loss = global_cvt_2_tf_float(find_atomic) * tf.reduce_mean( tf.square(self.scale*(polar - atomic_polar_hat)), name='l2_'+suffix)
+            more_loss['local_loss'] = local_loss
+            l2_loss += self.local_weight * local_loss
+            self.l2_loss_local_summary = tf.summary.scalar('l2_local_loss', 
+                                            tf.sqrt(more_loss['local_loss']))
+        
+
+        if self.global_weight > 0.0:    # Need global loss
+            atoms = 0
+            if self.type_sel is not None:
+                for w in self.type_sel:
+                    atoms += natoms[2+w]
+            else:
+                atoms = natoms[0]     
+            nframes = tf.shape(polar)[0] // self.tensor_size // atoms
+            # get global results
+            global_polar = tf.reshape(tf.reduce_sum(tf.reshape(
+                polar, [nframes, -1, self.tensor_size]), axis=1),[-1])
+            #if self.atomic: # If label is local, however
+            #    global_polar_hat = tf.reshape(tf.reduce_sum(tf.reshape(
+            #        polar_hat, [nframes, -1, self.tensor_size]), axis=1),[-1])
+            #else:
+            #    global_polar_hat = polar_hat
+            
+            global_loss = global_cvt_2_tf_float(find_global) * tf.reduce_mean( tf.square(self.scale*(global_polar - polar_hat)), name='l2_'+suffix)
+
+            more_loss['global_loss'] = global_loss
+            self.l2_loss_global_summary = tf.summary.scalar('l2_global_loss', 
+                                            tf.sqrt(more_loss['global_loss']) / global_cvt_2_tf_float(atoms))
+
+            # YWolfeee: should only consider atoms with dipole, i.e. atoms
+            # atom_norm  = 1./ global_cvt_2_tf_float(natoms[0])  
+            atom_norm  = 1./ global_cvt_2_tf_float(atoms)  
+            global_loss *= atom_norm   
+
+            l2_loss += self.global_weight * global_loss
+            
+        self.l2_more = more_loss
+        self.l2_l = l2_loss
+
+        self.l2_loss_summary = tf.summary.scalar('l2_loss', tf.sqrt(l2_loss))
+        return l2_loss, more_loss
+
+    def eval(self, sess, feed_dict, natoms):
+        atoms = 0
+        if self.type_sel is not None:
+            for w in self.type_sel:
+                atoms += natoms[2+w]
+        else:
+            atoms = natoms[0]
+
+        run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
+        error, error_lc, error_gl = run_sess(sess, run_data, feed_dict=feed_dict)
+
+        results = {"natoms": atoms, "rmse": np.sqrt(error)}
+        if self.local_weight > 0.0:
+            results["rmse_lc"] = np.sqrt(error_lc)
+        if self.global_weight > 0.0:
+            results["rmse_gl"] = np.sqrt(error_gl) / atoms
+        return results
+
+    def print_header(self):  # depreciated
+        prop_fmt = '   %11s %11s'
+        print_str = ''
+        print_str += prop_fmt % ('rmse_tst', 'rmse_trn')
+        if self.local_weight > 0.0:
+            print_str += prop_fmt % ('rmse_lc_tst', 'rmse_lc_trn')
+        if self.global_weight > 0.0:
+            print_str += prop_fmt % ('rmse_gl_tst', 'rmse_gl_trn')
+        return print_str
+
+    def print_on_training(self, 
+                          tb_writer,
+                          cur_batch,
+                          sess, 
+                          natoms,
+                          feed_dict_test,
+                          feed_dict_batch) :  # depreciated
+
+        # YHT: added to calculate the atoms number
+        atoms = 0
+        if self.type_sel is not None:
+            for w in self.type_sel:
+                atoms += natoms[2+w]                   
+        else:
+            atoms = natoms[0]
+
+        run_data = [self.l2_l, self.l2_more['local_loss'], self.l2_more['global_loss']]
+        summary_list = [self.l2_loss_summary]
+        if self.local_weight > 0.0:
+            summary_list.append(self.l2_loss_local_summary)
+        if self.global_weight > 0.0:
+            summary_list.append(self.l2_loss_global_summary)
+
+        # first train data
+        error_train = run_sess(sess, run_data, feed_dict=feed_dict_batch)
+
+        # than test data, if tensorboard log writter is present, commpute summary
+        # and write tensorboard logs
+        if tb_writer:
+            #summary_merged_op = tf.summary.merge([self.l2_loss_summary])
+            summary_merged_op = tf.summary.merge(summary_list)
+            run_data.insert(0, summary_merged_op)
+
+        test_out = run_sess(sess, run_data, feed_dict=feed_dict_test)
+
+        if tb_writer:
+            summary = test_out.pop(0)
+            tb_writer.add_summary(summary, cur_batch)
+
+        error_test = test_out  
+        
+        print_str = ""
+        prop_fmt = "   %11.2e %11.2e"
+        print_str += prop_fmt % (np.sqrt(error_test[0]), np.sqrt(error_train[0]))
+        if self.local_weight > 0.0:
+            print_str += prop_fmt % (np.sqrt(error_test[1]), np.sqrt(error_train[1]) )
+        if self.global_weight > 0.0:
+            print_str += prop_fmt % (np.sqrt(error_test[2])/atoms, np.sqrt(error_train[2])/atoms)
+
+        return print_str

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/model/__init__.py

+from .ener import EnerModel
+from .tensor import WFCModel
+from .tensor import DipoleModel
+from .tensor import PolarModel
+from .tensor import GlobalPolarModel

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/model/ener.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.utils.pair_tab import PairTab
+from deepmd.utils.graph import load_graph_def, get_tensor_by_name_from_graph
+from deepmd.utils.errors import GraphWithoutTensorError
+from deepmd.common import ClassArg
+from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import op_module
+from .model import Model
+from .model_stat import make_stat_input, merge_sys_stat
+
+class EnerModel(Model) :
+    """Energy model.
+    
+    Parameters
+    ----------
+    descrpt
+            Descriptor
+    fitting
+            Fitting net
+    type_map
+            Mapping atom type to the name (str) of the type.
+            For example `type_map[1]` gives the name of the type 1.
+    data_stat_nbatch
+            Number of frames used for data statistic
+    data_stat_protect
+            Protect parameter for atomic energy regression
+    use_srtab
+            The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.
+    smin_alpha
+            The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when `use_srtab` is provided.
+    sw_rmin
+            The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
+    sw_rmin
+            The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when `use_srtab` is provided.
+    """
+    model_type = 'ener'
+
+    def __init__ (
+            self, 
+            descrpt, 
+            fitting, 
+            typeebd = None,
+            type_map : List[str] = None,
+            data_stat_nbatch : int = 10,
+            data_stat_protect : float = 1e-2,
+            use_srtab : str = None,
+            smin_alpha : float = None,
+            sw_rmin : float = None,
+            sw_rmax : float = None
+    ) -> None:
+        """
+        Constructor
+        """
+        # descriptor
+        self.descrpt = descrpt
+        self.rcut = self.descrpt.get_rcut()
+        self.ntypes = self.descrpt.get_ntypes()
+        # fitting
+        self.fitting = fitting
+        self.numb_fparam = self.fitting.get_numb_fparam()
+        # type embedding
+        self.typeebd = typeebd
+        # other inputs
+        if type_map is None:
+            self.type_map = []
+        else:
+            self.type_map = type_map
+        self.data_stat_nbatch = data_stat_nbatch
+        self.data_stat_protect = data_stat_protect
+        self.srtab_name = use_srtab
+        if self.srtab_name is not None :
+            self.srtab = PairTab(self.srtab_name)
+            self.smin_alpha = smin_alpha
+            self.sw_rmin = sw_rmin
+            self.sw_rmax = sw_rmax
+        else :
+            self.srtab = None
+
+
+    def get_rcut (self) :
+        return self.rcut
+
+    def get_ntypes (self) :
+        return self.ntypes
+
+    def get_type_map (self) :
+        return self.type_map
+
+    def data_stat(self, data):
+        all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
+        m_all_stat = merge_sys_stat(all_stat)
+        self._compute_input_stat(m_all_stat, protection=self.data_stat_protect, mixed_type=data.mixed_type)
+        self._compute_output_stat(all_stat, mixed_type=data.mixed_type)
+        # self.bias_atom_e = data.compute_energy_shift(self.rcond)
+
+    def _compute_input_stat (self, all_stat, protection=1e-2, mixed_type=False):
+        if mixed_type:
+            self.descrpt.compute_input_stats(all_stat['coord'],
+                                             all_stat['box'],
+                                             all_stat['type'],
+                                             all_stat['natoms_vec'],
+                                             all_stat['default_mesh'],
+                                             all_stat,
+                                             mixed_type,
+                                             all_stat['real_natoms_vec'])
+        else:
+            self.descrpt.compute_input_stats(all_stat['coord'],
+                                             all_stat['box'],
+                                             all_stat['type'],
+                                             all_stat['natoms_vec'],
+                                             all_stat['default_mesh'],
+                                             all_stat)
+        self.fitting.compute_input_stats(all_stat, protection=protection)
+
+    def _compute_output_stat (self, all_stat, mixed_type=False):
+        if mixed_type:
+            self.fitting.compute_output_stats(all_stat, mixed_type=mixed_type)
+        else:
+            self.fitting.compute_output_stats(all_stat)
+
+
+    def build (self, 
+               coord_, 
+               atype_,
+               natoms,
+               box, 
+               mesh,
+               input_dict,
+               frz_model = None,
+               suffix = '', 
+               reuse = None):
+ 
+        if input_dict is None:
+            input_dict = {}
+        with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
+            t_tmap = tf.constant(' '.join(self.type_map), 
+                                 name = 'tmap', 
+                                 dtype = tf.string)
+            t_mt = tf.constant(self.model_type, 
+                               name = 'model_type', 
+                               dtype = tf.string)
+            t_ver = tf.constant(MODEL_VERSION,
+                                name = 'model_version',
+                                dtype = tf.string)
+
+            if self.srtab is not None :
+                tab_info, tab_data = self.srtab.get()
+                self.tab_info = tf.get_variable('t_tab_info',
+                                                tab_info.shape,
+                                                dtype = tf.float64,
+                                                trainable = False,
+                                                initializer = tf.constant_initializer(tab_info, dtype = tf.float64))
+                self.tab_data = tf.get_variable('t_tab_data',
+                                                tab_data.shape,
+                                                dtype = tf.float64,
+                                                trainable = False,
+                                                initializer = tf.constant_initializer(tab_data, dtype = tf.float64))
+
+        coord = tf.reshape (coord_, [-1, natoms[1] * 3])
+        atype = tf.reshape (atype_, [-1, natoms[1]])
+        input_dict['nframes'] = tf.shape(coord)[0]
+
+        # type embedding if any
+        if self.typeebd is not None:
+            type_embedding = self.typeebd.build(
+                self.ntypes,
+                reuse = reuse,
+                suffix = suffix,
+            )
+            input_dict['type_embedding'] = type_embedding
+            input_dict['atype'] = atype_
+
+        if frz_model == None:
+            dout \
+                = self.descrpt.build(coord_,
+                                     atype_,
+                                     natoms,
+                                     box,
+                                     mesh,
+                                     input_dict,
+                                     suffix = suffix,
+                                     reuse = reuse)
+            dout = tf.identity(dout, name='o_descriptor')
+        else:
+            tf.constant(self.rcut,
+                name = 'descrpt_attr/rcut',
+                dtype = GLOBAL_TF_FLOAT_PRECISION)
+            tf.constant(self.ntypes,
+                name = 'descrpt_attr/ntypes',
+                dtype = tf.int32)
+            feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
+            return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
+            imported_tensors \
+                = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
+            dout = imported_tensors[-1]
+            self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
+
+
+        if self.srtab is not None :
+            nlist, rij, sel_a, sel_r = self.descrpt.get_nlist()
+            nnei_a = np.cumsum(sel_a)[-1]
+            nnei_r = np.cumsum(sel_r)[-1]
+
+        atom_ener = self.fitting.build (dout, 
+                                        natoms, 
+                                        input_dict, 
+                                        reuse = reuse, 
+                                        suffix = suffix)
+        self.atom_ener = atom_ener
+
+        if self.srtab is not None :
+            sw_lambda, sw_deriv \
+                = op_module.soft_min_switch(atype, 
+                                            rij, 
+                                            nlist,
+                                            natoms,
+                                            sel_a = sel_a,
+                                            sel_r = sel_r,
+                                            alpha = self.smin_alpha,
+                                            rmin = self.sw_rmin,
+                                            rmax = self.sw_rmax)            
+            inv_sw_lambda = 1.0 - sw_lambda
+            # NOTICE:
+            # atom energy is not scaled, 
+            # force and virial are scaled
+            tab_atom_ener, tab_force, tab_atom_virial \
+                = op_module.pair_tab(self.tab_info,
+                                      self.tab_data,
+                                      atype,
+                                      rij,
+                                      nlist,
+                                      natoms,
+                                      sw_lambda,
+                                      sel_a = sel_a,
+                                      sel_r = sel_r)
+            energy_diff = tab_atom_ener - tf.reshape(atom_ener, [-1, natoms[0]])
+            tab_atom_ener = tf.reshape(sw_lambda, [-1]) * tf.reshape(tab_atom_ener, [-1])
+            atom_ener = tf.reshape(inv_sw_lambda, [-1]) * atom_ener
+            energy_raw = tab_atom_ener + atom_ener
+        else :
+            energy_raw = atom_ener
+
+        energy_raw = tf.reshape(energy_raw, [-1, natoms[0]], name = 'o_atom_energy'+suffix)
+        energy = tf.reduce_sum(global_cvt_2_ener_float(energy_raw), axis=1, name='o_energy'+suffix)
+
+        force, virial, atom_virial \
+            = self.descrpt.prod_force_virial (atom_ener, natoms)
+
+        if self.srtab is not None :
+            sw_force \
+                = op_module.soft_min_force(energy_diff, 
+                                           sw_deriv,
+                                           nlist, 
+                                           natoms,
+                                           n_a_sel = nnei_a,
+                                           n_r_sel = nnei_r)
+            force = force + sw_force + tab_force
+
+        force = tf.reshape (force, [-1, 3 * natoms[1]], name = "o_force"+suffix)
+
+        if self.srtab is not None :
+            sw_virial, sw_atom_virial \
+                = op_module.soft_min_virial (energy_diff,
+                                             sw_deriv,
+                                             rij,
+                                             nlist,
+                                             natoms,
+                                             n_a_sel = nnei_a,
+                                             n_r_sel = nnei_r)
+            atom_virial = atom_virial + sw_atom_virial + tab_atom_virial
+            virial = virial + sw_virial \
+                     + tf.reduce_sum(tf.reshape(tab_atom_virial, [-1, natoms[1], 9]), axis = 1)
+
+        virial = tf.reshape (virial, [-1, 9], name = "o_virial"+suffix)
+        atom_virial = tf.reshape (atom_virial, [-1, 9 * natoms[1]], name = "o_atom_virial"+suffix)
+
+        model_dict = {}
+        model_dict['energy'] = energy
+        model_dict['force'] = force
+        model_dict['virial'] = virial
+        model_dict['atom_ener'] = energy_raw
+        model_dict['atom_virial'] = atom_virial
+        model_dict['coord'] = coord
+        model_dict['atype'] = atype
+        
+        return model_dict
+
+    def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
+        graph, graph_def = load_graph_def(frz_model)
+        return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
+
+    def init_variables(self,
+                       graph : tf.Graph,
+                       graph_def : tf.GraphDef,
+                       model_type : str = "original_model",
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given frozen model
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        model_type : str
+            the type of the model
+        suffix : str
+            suffix to name scope
+        """
+        # self.frz_model will control the self.model to import the descriptor from the given frozen model instead of building from scratch...
+        # initialize fitting net with the given compressed frozen model
+        if model_type == 'original_model':
+            self.descrpt.init_variables(graph, graph_def, suffix=suffix)
+            self.fitting.init_variables(graph, graph_def, suffix=suffix)
+            tf.constant("original_model", name = 'model_type', dtype = tf.string)
+        elif model_type == 'compressed_model':
+            self.fitting.init_variables(graph, graph_def, suffix=suffix)
+            tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
+        else:
+            raise RuntimeError("Unknown model type %s" % model_type)
+        if self.typeebd is not None:
+            self.typeebd.init_variables(graph, graph_def, suffix=suffix)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/model/model.py

+from deepmd.env import tf
+
+
+class Model:
+    def init_variables(self,
+                       graph : tf.Graph,
+                       graph_def : tf.GraphDef,
+                       model_type : str = "original_model",
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given frozen model
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        model_type : str
+            the type of the model
+        suffix : str
+            suffix to name scope
+        """
+        raise RuntimeError("The 'dp train init-frz-model' command do not support this model!")

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/model/model_stat.py

+import numpy as np
+from collections import defaultdict
+
+def _make_all_stat_ref(data, nbatches):
+    all_stat = defaultdict(list)
+    for ii in range(data.get_nsystems()) :
+        for jj in range(nbatches) :
+            stat_data = data.get_batch (sys_idx = ii)
+            for dd in stat_data:
+                if dd == "natoms_vec":
+                    stat_data[dd] = stat_data[dd].astype(np.int32) 
+                all_stat[dd].append(stat_data[dd])        
+    return all_stat
+
+
+def make_stat_input(data, nbatches, merge_sys = True):
+    """
+    pack data for statistics
+
+    Parameters
+    ----------
+    data:
+        The data
+    merge_sys: bool (True)
+        Merge system data
+
+    Returns
+    -------
+    all_stat:
+        A dictionary of list of list storing data for stat. 
+        if merge_sys == False data can be accessed by 
+            all_stat[key][sys_idx][batch_idx][frame_idx]
+        else merge_sys == True can be accessed by 
+            all_stat[key][batch_idx][frame_idx]
+    """
+    all_stat = defaultdict(list)
+    for ii in range(data.get_nsystems()) :
+        sys_stat =  defaultdict(list)
+        for jj in range(nbatches) :
+            stat_data = data.get_batch (sys_idx = ii)
+            for dd in stat_data:
+                if dd == "natoms_vec":
+                    stat_data[dd] = stat_data[dd].astype(np.int32) 
+                sys_stat[dd].append(stat_data[dd])
+        for dd in sys_stat:
+            if merge_sys:
+                for bb in sys_stat[dd]:
+                    all_stat[dd].append(bb)
+            else:                    
+                all_stat[dd].append(sys_stat[dd])
+    return all_stat
+
+def merge_sys_stat(all_stat):
+    first_key = list(all_stat.keys())[0]
+    nsys = len(all_stat[first_key])
+    ret = defaultdict(list)
+    for ii in range(nsys):
+        for dd in all_stat:
+            for bb in all_stat[dd][ii]:
+                ret[dd].append(bb)
+    return ret
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/model/tensor.py

+import numpy as np
+from typing import Tuple, List
+
+from deepmd.env import tf
+from deepmd.common import ClassArg
+from deepmd.env import global_cvt_2_ener_float, MODEL_VERSION, GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.utils.graph import load_graph_def
+from .model import Model
+from .model_stat import make_stat_input, merge_sys_stat
+
+class TensorModel(Model) :
+    """Tensor model.
+
+    Parameters
+    ----------
+    tensor_name
+            Name of the tensor.
+    descrpt
+            Descriptor
+    fitting
+            Fitting net
+    type_map
+            Mapping atom type to the name (str) of the type.
+            For example `type_map[1]` gives the name of the type 1.
+    data_stat_nbatch
+            Number of frames used for data statistic
+    data_stat_protect
+            Protect parameter for atomic energy regression
+    """
+    def __init__ (
+            self, 
+            tensor_name : str,
+            descrpt, 
+            fitting, 
+            type_map : List[str] = None,
+            data_stat_nbatch : int = 10,
+            data_stat_protect : float = 1e-2,
+    )->None:
+        """
+        Constructor
+        """
+        self.model_type = tensor_name
+        # descriptor
+        self.descrpt = descrpt
+        self.rcut = self.descrpt.get_rcut()
+        self.ntypes = self.descrpt.get_ntypes()
+        # fitting
+        self.fitting = fitting
+        # other params
+        if type_map is None:
+            self.type_map = []
+        else:
+            self.type_map = type_map
+        self.data_stat_nbatch = data_stat_nbatch
+        self.data_stat_protect = data_stat_protect
+    
+    def get_rcut (self) :
+        return self.rcut
+
+    def get_ntypes (self) :
+        return self.ntypes
+
+    def get_type_map (self) :
+        return self.type_map
+
+    def get_sel_type(self):
+        return self.fitting.get_sel_type()
+
+    def get_out_size (self) :
+        return self.fitting.get_out_size()
+
+    def data_stat(self, data):
+        all_stat = make_stat_input(data, self.data_stat_nbatch, merge_sys = False)
+        m_all_stat = merge_sys_stat(all_stat)        
+        self._compute_input_stat (m_all_stat, protection = self.data_stat_protect)
+        self._compute_output_stat(all_stat)
+
+    def _compute_input_stat(self, all_stat, protection = 1e-2) :
+        self.descrpt.compute_input_stats(all_stat['coord'],
+                                         all_stat['box'],
+                                         all_stat['type'],
+                                         all_stat['natoms_vec'],
+                                         all_stat['default_mesh'], 
+                                         all_stat)
+        if hasattr(self.fitting, 'compute_input_stats'):
+            self.fitting.compute_input_stats(all_stat, protection = protection)
+
+    def _compute_output_stat (self, all_stat) :
+        if hasattr(self.fitting, 'compute_output_stats'):
+            self.fitting.compute_output_stats(all_stat)
+
+    def build (self, 
+               coord_, 
+               atype_,
+               natoms,
+               box, 
+               mesh,
+               input_dict,
+               frz_model = None,         
+               suffix = '', 
+               reuse = None):
+        with tf.variable_scope('model_attr' + suffix, reuse = reuse) :
+            t_tmap = tf.constant(' '.join(self.type_map), 
+                                 name = 'tmap', 
+                                 dtype = tf.string)
+            t_st = tf.constant(self.get_sel_type(), 
+                               name = 'sel_type',
+                               dtype = tf.int32)
+            t_mt = tf.constant(self.model_type, 
+                               name = 'model_type', 
+                               dtype = tf.string)
+            t_ver = tf.constant(MODEL_VERSION,
+                                name = 'model_version',
+                                dtype = tf.string)
+            t_od = tf.constant(self.get_out_size(), 
+                               name = 'output_dim', 
+                               dtype = tf.int32)
+
+        natomsel = sum(natoms[2+type_i] for type_i in self.get_sel_type())
+        nout = self.get_out_size()
+
+        if frz_model == None:
+            dout \
+                = self.descrpt.build(coord_,
+                                     atype_,
+                                     natoms,
+                                     box,
+                                     mesh,
+                                     input_dict,
+                                     suffix = suffix,
+                                     reuse = reuse)
+            dout = tf.identity(dout, name='o_descriptor')
+        else:
+            tf.constant(self.rcut,
+                name = 'descrpt_attr/rcut',
+                dtype = GLOBAL_TF_FLOAT_PRECISION)
+            tf.constant(self.ntypes,
+                name = 'descrpt_attr/ntypes',
+                dtype = tf.int32)
+            feed_dict = self.descrpt.get_feed_dict(coord_, atype_, natoms, box, mesh)
+            return_elements = [*self.descrpt.get_tensor_names(), 'o_descriptor:0']
+            imported_tensors \
+                = self._import_graph_def_from_frz_model(frz_model, feed_dict, return_elements)
+            dout = imported_tensors[-1]
+            self.descrpt.pass_tensors_from_frz_model(*imported_tensors[:-1])
+
+        rot_mat = self.descrpt.get_rot_mat()
+        rot_mat = tf.identity(rot_mat, name = 'o_rot_mat'+suffix)
+
+        output = self.fitting.build (dout, 
+                                     rot_mat,
+                                     natoms, 
+                                     reuse = reuse, 
+                                     suffix = suffix)
+        framesize = nout if "global" in self.model_type else natomsel * nout
+        output = tf.reshape(output, [-1, framesize], name = 'o_' + self.model_type + suffix)
+
+        model_dict = {self.model_type: output}
+
+        if "global" not in self.model_type:
+            gname = "global_"+self.model_type
+            atom_out = tf.reshape(output, [-1, natomsel, nout])
+            global_out = tf.reduce_sum(atom_out, axis=1)
+            global_out = tf.reshape(global_out, [-1, nout], name="o_" + gname + suffix)
+            
+            out_cpnts = tf.split(atom_out, nout, axis=-1)
+            force_cpnts = []
+            virial_cpnts = []
+            atom_virial_cpnts = []
+
+            for out_i in out_cpnts:
+                force_i, virial_i, atom_virial_i \
+                    = self.descrpt.prod_force_virial(out_i, natoms)
+                force_cpnts.append      (tf.reshape(force_i,       [-1, 3*natoms[1]]))
+                virial_cpnts.append     (tf.reshape(virial_i,      [-1, 9]))
+                atom_virial_cpnts.append(tf.reshape(atom_virial_i, [-1, 9*natoms[1]]))
+
+            # [nframe x nout x (natom x 3)]
+            force = tf.concat(force_cpnts, axis=1, name="o_force" + suffix)
+            # [nframe x nout x 9]
+            virial = tf.concat(virial_cpnts, axis=1, name="o_virial" + suffix)
+            # [nframe x nout x (natom x 9)]
+            atom_virial = tf.concat(atom_virial_cpnts, axis=1, name="o_atom_virial" + suffix)
+
+            model_dict[gname] = global_out
+            model_dict["force"] = force
+            model_dict["virial"] = virial
+            model_dict["atom_virial"] = atom_virial
+
+        return model_dict
+
+    def _import_graph_def_from_frz_model(self, frz_model, feed_dict, return_elements):
+        graph, graph_def = load_graph_def(frz_model)
+        return tf.import_graph_def(graph_def, input_map = feed_dict, return_elements = return_elements, name = "")
+
+    def init_variables(self,
+                       graph : tf.Graph,
+                       graph_def : tf.GraphDef,
+                       model_type : str = "original_model",
+                       suffix : str = "",
+    ) -> None:
+        """
+        Init the embedding net variables with the given frozen model
+
+        Parameters
+        ----------
+        graph : tf.Graph
+            The input frozen model graph
+        graph_def : tf.GraphDef
+            The input frozen model graph_def
+        model_type : str
+            the type of the model
+        suffix : str
+            suffix to name scope
+        """
+        if model_type == 'original_model':
+            self.descrpt.init_variables(graph, graph_def, suffix=suffix)
+            self.fitting.init_variables(graph, graph_def, suffix=suffix)
+            tf.constant("original_model", name = 'model_type', dtype = tf.string)
+        elif model_type == 'compressed_model':
+            self.fitting.init_variables(graph, graph_def, suffix=suffix)
+            tf.constant("compressed_model", name = 'model_type', dtype = tf.string)
+        else:
+            raise RuntimeError("Unknown model type %s" % model_type)
+
+
+class WFCModel(TensorModel):
+    def __init__(
+            self, 
+            descrpt, 
+            fitting, 
+            type_map : List[str] = None, 
+            data_stat_nbatch : int = 10, 
+            data_stat_protect : float = 1e-2
+    ) -> None:
+        TensorModel.__init__(self, 'wfc', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
+
+class DipoleModel(TensorModel):
+    def __init__(
+            self, 
+            descrpt, 
+            fitting, 
+            type_map : List[str] = None, 
+            data_stat_nbatch : int = 10, 
+            data_stat_protect : float = 1e-2
+    ) -> None:
+        TensorModel.__init__(self, 'dipole', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
+
+class PolarModel(TensorModel):
+    def __init__(
+            self, 
+            descrpt, 
+            fitting, 
+            type_map : List[str] = None, 
+            data_stat_nbatch : int = 10, 
+            data_stat_protect : float = 1e-2
+    ) -> None:
+        TensorModel.__init__(self, 'polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
+
+class GlobalPolarModel(TensorModel):
+    def __init__(
+            self, 
+            descrpt, 
+            fitting, 
+            type_map : List[str] = None, 
+            data_stat_nbatch : int = 10, 
+            data_stat_protect : float = 1e-2
+    ) -> None:
+        TensorModel.__init__(self, 'global_polar', descrpt, fitting, type_map, data_stat_nbatch, data_stat_protect)
+
+

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/nvnmd/__init__.py

+
+from . import data, descriptor, entrypoints, fit, utils
+
+__all__ = [
+    "data",
+    "descriptor",
+    "entrypoints",
+    "fit",
+    "utils",
+]

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/nvnmd/data/__init__.py

+"""
+nvnmd.data
+==========
+
+Provides
+    1. hardware configuration
+    2. default input script
+    3. title and citation
+
+Data
+----
+
+jdata_sys 
+    action configuration
+jdata_config 
+    hardware configuration
+
+    dscp 
+        descriptor configuration
+    fitn 
+        fitting network configuration
+    size 
+        ram capacity
+    ctrl 
+        control flag, such as Time Division Multiplexing (TDM)
+    nbit 
+        number of bits of fixed-point number
+jdata_config_16 (disable) 
+    difference with configure fitting size as 16
+jdata_config_32 (disable) 
+    difference with configure fitting size as 32
+jdata_config_64 (disable) 
+    difference with configure fitting size as 64
+jdata_config_128 (default) 
+    difference with configure fitting size as 128
+jdata_configs 
+    all configure of jdata_config{nfit_node}
+jdata_deepmd_input 
+    default input script for nvnmd training
+NVNMD_WELCOME 
+    nvnmd title when logging
+NVNMD_CITATION 
+    citation of nvnmd
+"""

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/nvnmd/data/data.py

+
+jdata_sys = {
+    "debug": False
+}
+
+jdata_config = {
+    "dscp": {
+        "sel": [60, 60],
+        "rcut": 6.0,
+        "rcut_smth": 0.5,
+        "neuron": [8, 16, 32],
+        "resnet_dt": False,
+        "axis_neuron": 4,
+        "type_one_side": True,
+
+        "NI": 128,
+        "rc_lim": 0.5,
+        "M1": "neuron[-1]",
+        "M2": "axis_neuron",
+        "SEL": [60, 60, 0, 0],
+        "NNODE_FEAS": "(1, neuron)",
+        "nlayer_fea": "len(neuron)",
+        "same_net": "type_one_side",
+        "NIDP": "sum(sel)",
+        "NIX": "2^ceil(ln2(NIDP/1.5))",
+        "ntype": "len(sel)",
+        "ntypex": "same_net ? 1: ntype",
+        "ntypex_max": 1,
+        "ntype_max": 4
+    },
+
+    "fitn": {
+        "neuron": [32, 32, 32],
+        "resnet_dt": False,
+
+        "NNODE_FITS": "(M1*M2, neuron, 1)",
+        "nlayer_fit": "len(neuron)+1",
+        "NLAYER": "nlayer_fit"
+    },
+
+    "size": {
+        "NTYPE_MAX": 4,
+        "NSPU": 4096,
+        "MSPU": 32768,
+        "Na": "NSPU",
+        "NaX": "MSPU"
+    },
+
+    "ctrl": {
+        "NSTDM": 16,
+        "NSTDM_M1": 16,
+        "NSTDM_M2": 1,
+        "NSADV": "NSTDM+1",
+        "NSEL": "NSTDM*ntype_max",
+        "NSTDM_M1X": 4,
+        "NSTEP_DELAY": 20,
+        "MAX_FANOUT": 30
+    },
+
+    "nbit": {
+        "NBIT_DATA": 21,
+        "NBIT_DATA_FL": 13,
+        "NBIT_LONG_DATA": 32,
+        "NBIT_LONG_DATA_FL": 24,
+        "NBIT_DIFF_DATA": 24,
+
+        "NBIT_SPE": 2,
+        "NBIT_CRD": "NBIT_DATA*3",
+        "NBIT_LST": "ln2(NaX)",
+
+        "NBIT_SPE_MAX": 8,
+        "NBIT_LST_MAX": 16,
+
+        "NBIT_ATOM": "NBIT_SPE+NBIT_CRD",
+        "NBIT_LONG_ATOM": "NBIT_SPE+NBIT_LONG_DATA*3",
+
+        "NBIT_RIJ": "NBIT_DATA_FL+5",
+        "NBIT_FEA_X": 10,
+        "NBIT_FEA_X_FL": 4,
+        "NBIT_FEA_X2_FL": 6,
+        "NBIT_FEA": 18,
+        "NBIT_FEA_FL": 10,
+        "NBIT_SHIFT": 4,
+
+        "NBIT_DATA2": "NBIT_DATA+NBIT_DATA_FL",
+        "NBIT_DATA2_FL": "2*NBIT_DATA_FL",
+        "NBIT_DATA_FEA": "NBIT_DATA+NBIT_FEA_FL",
+        "NBIT_DATA_FEA_FL": "NBIT_DATA_FL+NBIT_FEA_FL",
+
+        "NBIT_FORCE": 32,
+        "NBIT_FORCE_FL": "2*NBIT_DATA_FL-1",
+
+        "NBIT_SUM": "NBIT_DATA_FL+8",
+        "NBIT_WEIGHT": 18,
+        "NBIT_WEIGHT_FL": 13,
+
+        "NBIT_RAM": 72,
+        "NBIT_ADDR": 32,
+
+        "NBTI_MODEL_HEAD": 32,
+
+        "NBIT_TH_LONG_ADD": 30,
+        "NBIT_ADD": 15,
+
+        "RANGE_B": [-100, 100],
+        "RANGE_W": [-20, 20],
+
+        "NCFG": 35,
+        "NNET": 4920,
+        "NFEA": 8192
+    },
+
+    "end": ""
+}
+
+jdata_config_16 = {
+    "dscp": {
+        "neuron": [8, 16, 32],
+        "axis_neuron": 4,
+        "NI": 128
+    },
+
+    "fitn": {
+        "neuron": [16, 16, 16]
+    },
+
+    "ctrl": {
+        "NSTDM": 16,
+        "NSTDM_M1": 16,
+        "NSTDM_M2": 1,
+        "NSTDM_M1X": 4
+    }
+}
+
+jdata_config_32 = {
+    "dscp": {
+        "neuron": [8, 16, 32],
+        "axis_neuron": 4,
+        "NI": 128
+    },
+
+    "fitn": {
+        "neuron": [32, 32, 32]
+    },
+
+    "ctrl": {
+        "NSTDM": 16,
+        "NSTDM_M1": 16,
+        "NSTDM_M2": 1,
+        "NSTDM_M1X": 4
+    }
+}
+
+jdata_config_64 = {
+    "dscp": {
+        "neuron": [8, 16, 32],
+        "axis_neuron": 4,
+        "NI": 128
+    },
+
+    "fitn": {
+        "neuron": [64, 64, 64]
+    },
+
+    "ctrl": {
+        "NSTDM": 32,
+        "NSTDM_M1": 32,
+        "NSTDM_M2": 1,
+        "NSTDM_M1X": 4
+    }
+}
+
+jdata_config_128 = {
+    "dscp": {
+        "neuron": [8, 16, 32],
+        "axis_neuron": 4,
+        "NI": 128
+    },
+
+    "fitn": {
+        "neuron": [128, 128, 128]
+    },
+
+    "ctrl": {
+        "NSTDM": 32,
+        "NSTDM_M1": 32,
+        "NSTDM_M2": 1,
+        "NSTDM_M1X": 4
+    }
+}
+
+jdata_configs = {
+    "_16": jdata_config_16,
+    "_32": jdata_config_32,
+    "_64": jdata_config_64,
+    "128": jdata_config_128
+}
+
+jdata_deepmd_input = {
+    "model": {
+        "descriptor": {
+            "seed": 1,
+            "type": "se_a",
+            "sel": [
+                60,
+                60
+            ],
+            "rcut": 7.0,
+            "rcut_smth": 0.5,
+            "neuron": [
+                8,
+                16,
+                32
+            ],
+            "type_one_side": False,
+            "axis_neuron": 4,
+            "resnet_dt": False
+        },
+        "fitting_net": {
+            "seed": 1,
+            "neuron": [
+                128,
+                128,
+                128
+            ],
+            "resnet_dt": False
+        }
+    },
+    "nvnmd": {
+        "net_size": 128,
+        "config_file": "none",
+        "weight_file": "none",
+        "map_file": "none",
+        "enable": False,
+        "restore_descriptor": False,
+        "restore_fitting_net": False,
+        "quantize_descriptor": False,
+        "quantize_fitting_net": False
+    },
+    "learning_rate": {
+        "type": "exp",
+        "decay_steps": 5000,
+        "start_lr": 0.005,
+        "stop_lr": 8.257687192506788e-05
+    },
+    "loss": {
+        "start_pref_e": 0.02,
+        "limit_pref_e": 1,
+        "start_pref_f": 1000,
+        "limit_pref_f": 1,
+        "start_pref_v": 0,
+        "limit_pref_v": 0
+    },
+    "training": {
+        "seed": 1,
+        "stop_batch": 10000,
+        "disp_file": "lcurve.out",
+        "disp_freq": 100,
+        "numb_test": 10,
+        "save_freq": 1000,
+        "save_ckpt": "model.ckpt",
+        "disp_training": True,
+        "time_training": True,
+        "profiling": False,
+        "training_data": {
+            "systems": "dataset",
+            "set_prefix": "set",
+            "batch_size": 1
+        }
+    }
+}
+NVNMD_WELCOME = (
+    " _   _  __     __  _   _   __  __   ____  ",
+    "| \ | | \ \   / / | \ | | |  \/  | |  _ \ ",
+    "|  \| |  \ \ / /  |  \| | | |\/| | | | | |",
+    "| |\  |   \ V /   | |\  | | |  | | | |_| |",
+    "|_| \_|    \_/    |_| \_| |_|  |_| |____/ ",
+)
+
+NVNMD_CITATION = (
+    "Please read and cite:",
+    "Mo et al., npj Comput Mater 8, 107 (2022)",
+)

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/nvnmd/descriptor/__init__.py

+"""
+nvnmd.se_a
+==========
+
+Provides
+    1. building descriptor with continuous embedding network
+    2. building descriptor with quantized embedding network
+
+"""

_skbuild/linux-x86_64-3.6/cmake-install/deepmd/nvnmd/descriptor/se_a.py

+import numpy as np
+
+from deepmd.env import tf
+from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
+from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
+from deepmd.env import op_module
+from deepmd.utils.network import embedding_net
+
+
+#
+from deepmd.nvnmd.utils.config import nvnmd_cfg
+from deepmd.nvnmd.utils.network import matmul3_qq
+from deepmd.nvnmd.utils.weight import get_normalize, get_rng_s
+
+
+def build_davg_dstd():
+    r"""Get the davg and dstd from the dictionary nvnmd_cfg.
+    The davg and dstd have been obtained by training CNN
+    """
+    davg, dstd = get_normalize(nvnmd_cfg.weight)
+    return davg, dstd
+
+
+def build_op_descriptor():
+    r"""Replace se_a.py/DescrptSeA/build
+    """
+    if nvnmd_cfg.quantize_descriptor:
+        return op_module.prod_env_mat_a_nvnmd_quantize
+    else:
+        return op_module.prod_env_mat_a
+
+
+def descrpt2r4(inputs, natoms):
+    r"""Replace :math:`r_{ji} \rightarrow r'_{ji}`
+    where :math:`r_{ji} = (x_{ji}, y_{ji}, z_{ji})` and
+    :math:`r'_{ji} = (s_{ji}, \frac{s_{ji} x_{ji}}{r_{ji}}, \frac{s_{ji} y_{ji}}{r_{ji}}, \frac{s_{ji} z_{ji}}{r_{ji}})`
+    """
+    NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
+    NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
+    NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
+    prec = 1.0 / (2 ** NBIT_FEA_X_FL)
+
+    ntypes = nvnmd_cfg.dscp['ntype']
+    NIDP = nvnmd_cfg.dscp['NIDP']
+    ndescrpt = NIDP * 4
+    start_index = 0
+
+    # (nf, na*nd)
+    shape = inputs.get_shape().as_list()
+    # (nf*na*ni, 4)
+    inputs_reshape = tf.reshape(inputs, [-1, 4])
+
+    with tf.variable_scope('filter_type_all_x', reuse=True):
+        # u (i.e., r^2)
+        u = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
+        with tf.variable_scope('u', reuse=True):
+            u = op_module.quantize_nvnmd(u, 0, -1, NBIT_DATA_FL, -1)
+        # print('u:', u)
+        u = tf.reshape(u, [-1, natoms[0] * NIDP])
+        # rij
+        rij = tf.reshape(tf.slice(inputs_reshape, [0, 1], [-1, 3]), [-1, 3])
+        with tf.variable_scope('rij', reuse=True):
+            rij = op_module.quantize_nvnmd(rij, 0, NBIT_DATA_FL, -1, -1)
+        # print('rij:', rij)
+        s = []
+        sr = []
+        for type_i in range(ntypes):
+            type_input = 0
+            postfix = f"_t{type_input}_t{type_i}"
+            u_i = tf.slice(
+                u,
+                [0, start_index * NIDP],
+                [-1, natoms[2 + type_i] * NIDP])
+            u_i = tf.reshape(u_i, [-1, 1])
+            #
+            keys = 's,sr'.split(',')
+            map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
+            map_tables2 = [nvnmd_cfg.map[f"d{key}_dr2" + postfix] for key in keys]
+            map_outs = []
+            for ii in range(len(keys)):
+                map_outs.append(op_module.map_nvnmd(
+                    u_i,
+                    map_tables[ii][0],
+                    map_tables[ii][1] / prec,
+                    map_tables2[ii][0],
+                    map_tables2[ii][1] / prec,
+                    prec, NBIT_FEA_FL))
+
+            s_i, sr_i = map_outs
+            s_i = tf.reshape(s_i, [-1, natoms[2 + type_i] * NIDP])
+            sr_i = tf.reshape(sr_i, [-1, natoms[2 + type_i] * NIDP])
+            s.append(s_i)
+            sr.append(sr_i)
+            start_index += natoms[2 + type_i]
+
+        s = tf.concat(s, axis=1)
+        sr = tf.concat(sr, axis=1)
+
+        with tf.variable_scope('s', reuse=True):
+            s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
+
+        with tf.variable_scope('sr', reuse=True):
+            sr = op_module.quantize_nvnmd(sr, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
+
+        s = tf.reshape(s, [-1, 1])
+        sr = tf.reshape(sr, [-1, 1])
+
+        # R2R4
+        Rs = s
+        Rxyz = sr * rij
+        with tf.variable_scope('Rxyz', reuse=True):
+            Rxyz = op_module.quantize_nvnmd(Rxyz, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
+        R4 = tf.concat([Rs, Rxyz], axis=1)
+        R4 = tf.reshape(R4, [-1, NIDP, 4])
+        inputs_reshape = R4
+        inputs_reshape = tf.reshape(inputs_reshape, [-1, ndescrpt])
+    return inputs_reshape
+
+
+def filter_lower_R42GR(
+        type_i,
+        type_input,
+        inputs_i,
+        is_exclude,
+        activation_fn,
+        bavg,
+        stddev,
+        trainable,
+        suffix,
+        seed,
+        seed_shift,
+        uniform_seed,
+        filter_neuron,
+        filter_precision,
+        filter_resnet_dt,
+        embedding_net_variables):
+    r"""Replace se_a.py/DescrptSeA/_filter_lower
+    """
+    shape_i = inputs_i.get_shape().as_list()
+    inputs_reshape = tf.reshape(inputs_i, [-1, 4])
+    natom = tf.shape(inputs_i)[0]
+    M1 = nvnmd_cfg.dscp['M1']
+
+    NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
+    NBIT_FEA_X_FL = nvnmd_cfg.nbit['NBIT_FEA_X_FL']
+    NBIT_FEA_X2_FL = nvnmd_cfg.nbit['NBIT_FEA_X2_FL']
+    NBIT_FEA_FL = nvnmd_cfg.nbit['NBIT_FEA_FL']
+    prec = 1.0 / (2 ** NBIT_FEA_X2_FL)
+    type_input = 0 if (type_input < 0) else type_input
+    postfix = f"_t{type_input}_t{type_i}"
+
+    if (nvnmd_cfg.quantize_descriptor):
+        s_min, smax = get_rng_s(nvnmd_cfg.weight)
+        s_min = -2.0
+        # s_min = np.floor(s_min)
+        s = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
+        s = op_module.quantize_nvnmd(s, 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
+        # G
+        keys = 'G'.split(',')
+        map_tables = [nvnmd_cfg.map[key + postfix] for key in keys]
+        map_tables2 = [nvnmd_cfg.map[f"d{key}_ds" + postfix] for key in keys]
+        map_outs = []
+        for ii in range(len(keys)):
+            with tf.variable_scope(keys[ii], reuse=True):
+                map_outs.append(op_module.map_nvnmd(
+                    s - s_min,
+                    map_tables[ii][0], map_tables[ii][1] / prec,
+                    map_tables2[ii][0], map_tables2[ii][1] / prec,
+                    prec, NBIT_FEA_FL))
+                map_outs[ii] = op_module.quantize_nvnmd(map_outs[ii], 0, NBIT_FEA_FL, NBIT_DATA_FL, -1)
+        G = map_outs
+        # G
+        xyz_scatter = G
+        xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
+        # GR
+        inputs_reshape = tf.reshape(inputs_reshape, [-1, shape_i[1] // 4, 4])
+        GR = matmul3_qq(tf.transpose(inputs_reshape, [0, 2, 1]), xyz_scatter, -1)
+        GR = tf.reshape(GR, [-1, 4 * M1])
+        return GR
+
+    else:
+        xyz_scatter = tf.reshape(tf.slice(inputs_reshape, [0, 0], [-1, 1]), [-1, 1])
+        if nvnmd_cfg.restore_descriptor:
+            trainable = False
+            embedding_net_variables = {}
+            for key in nvnmd_cfg.weight.keys():
+                if 'filter_type' in key:
+                    key2 = key.replace('.', '/')
+                    embedding_net_variables[key2] = nvnmd_cfg.weight[key]
+
+        if (not is_exclude):
+            xyz_scatter = embedding_net(
+                xyz_scatter,
+                filter_neuron,
+                filter_precision,
+                activation_fn=activation_fn,
+                resnet_dt=filter_resnet_dt,
+                name_suffix=suffix,
+                stddev=stddev,
+                bavg=bavg,
+                seed=seed,
+                trainable=trainable,
+                uniform_seed=uniform_seed,
+                initial_variables=embedding_net_variables)
+            if (not uniform_seed) and (seed is not None):
+                seed += seed_shift
+        else:
+            # we can safely return the final xyz_scatter filled with zero directly
+            return tf.cast(tf.fill((natom, 4, M1), 0.), GLOBAL_TF_FLOAT_PRECISION)
+        # natom x nei_type_i x out_size
+        xyz_scatter = tf.reshape(xyz_scatter, (-1, shape_i[1] // 4, M1))
+        # When using tf.reshape(inputs_i, [-1, shape_i[1]//4, 4]) below
+        # [588 24] -> [588 6 4] correct
+        # but if sel is zero
+        # [588 0] -> [147 0 4] incorrect; the correct one is [588 0 4]
+        # So we need to explicitly assign the shape to tf.shape(inputs_i)[0] instead of -1
+        return tf.matmul(tf.reshape(inputs_i, [natom, shape_i[1] // 4, 4]), xyz_scatter, transpose_a=True)
+
+
+def filter_GR2D(xyz_scatter_1):
+    r"""Replace se_a.py/_filter
+    """
+    NIX = nvnmd_cfg.dscp['NIX']
+    NBIT_DATA_FL = nvnmd_cfg.nbit['NBIT_DATA_FL']
+    M1 = nvnmd_cfg.dscp['M1']
+    M2 = nvnmd_cfg.dscp['M2']
+
+    if (nvnmd_cfg.quantize_descriptor):
+        xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4 * M1])
+        # fix the number of bits of gradient
+        xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, -1, NBIT_DATA_FL, -1)
+        xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
+        with tf.variable_scope('GR', reuse=True):
+            xyz_scatter_1 = op_module.quantize_nvnmd(xyz_scatter_1, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
+        xyz_scatter_1 = tf.reshape(xyz_scatter_1, [-1, 4, M1])
+
+        # natom x 4 x outputs_size_2
+        xyz_scatter_2 = xyz_scatter_1
+        # natom x 3 x outputs_size_1
+        qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
+        # natom x outputs_size_2 x 3
+        qmat = tf.transpose(qmat, perm=[0, 2, 1])
+        # D': natom x outputs_size x outputs_size_2
+        result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
+        # D': natom x (outputs_size x outputs_size_2)
+        result = tf.reshape(result, [-1, M1 * M1])
+        #
+        index_subset = []
+        for ii in range(M1):
+            for jj in range(ii, ii + M2):
+                index_subset.append((ii * M1) + (jj % M1))
+        index_subset = tf.constant(np.int32(np.array(index_subset)))
+        result = tf.gather(result, index_subset, axis=1)
+
+        with tf.variable_scope('d', reuse=True):
+            result = op_module.quantize_nvnmd(result, 0, NBIT_DATA_FL, NBIT_DATA_FL, -1)
+    else:
+        # natom x 4 x outputs_size
+        xyz_scatter_1 = xyz_scatter_1 * (1.0 / NIX)
+        # natom x 4 x outputs_size_2
+        # xyz_scatter_2 = tf.slice(xyz_scatter_1, [0,0,0],[-1,-1,outputs_size_2])
+        xyz_scatter_2 = xyz_scatter_1
+        # natom x 3 x outputs_size_1
+        qmat = tf.slice(xyz_scatter_1, [0, 1, 0], [-1, 3, -1])
+        # natom x outputs_size_1 x 3
+        qmat = tf.transpose(qmat, perm=[0, 2, 1])
+        # natom x outputs_size x outputs_size_2
+        result = tf.matmul(xyz_scatter_1, xyz_scatter_2, transpose_a=True)
+        # natom x (outputs_size x outputs_size_2)
+        # result = tf.reshape(result, [-1, outputs_size_2 * outputs_size[-1]])
+        result = tf.reshape(result, [-1, M1 * M1])
+        #
+        index_subset = []
+        for ii in range(M1):
+            for jj in range(ii, ii + M2):
+                index_subset.append((ii * M1) + (jj % M1))
+        index_subset = tf.constant(np.int32(np.array(index_subset)))
+        result = tf.gather(result, index_subset, axis=1)
+
+    return result, qmat