make deep tensor deriv again

.gitignore

@@ -29,3 +29,5 @@ _templates
 API_CC
 dp/
 build_lammps/
+build_tests/
+build_cc_tests

deepmd/.gitignore

@@ -0,0 +1,3 @@
+op/_*.py
+pkg_config
+!op/__init__.py

deepmd/entrypoints/freeze.py

@@ -66,6 +66,10 @@ def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> Li
     elif model_type == "dipole":
         nodes += [
             "o_dipole",
+            "o_global_dipole",
+            "o_force",
+            "o_virial",
+            "o_atom_virial",
             "o_rmat",
             "o_rmat_deriv",
             "o_nlist",
@@ -78,6 +82,10 @@ def _make_node_names(model_type: str, modifier_type: Optional[str] = None) -> Li
     elif model_type == "polar":
         nodes += [
             "o_polar",
+            "o_global_polar",
+            "o_force",
+            "o_virial",
+            "o_atom_virial",
             "model_attr/sel_type",
             "model_attr/output_dim",
         ]

deepmd/infer/deep_eval.py

@@ -215,8 +215,9 @@ def reverse_map(vec : np.ndarray, imap : List[int]) -> np.ndarray:
                 Reverse mapped vector.
         """
         ret = np.zeros(vec.shape)        
-        for idx,ii in enumerate(imap) :
-            ret[:,ii,:] = vec[:,idx,:]
+        # for idx,ii in enumerate(imap) :
+        #     ret[:,ii,:] = vec[:,idx,:]
+        ret[:, imap, :] = vec
         return ret
 
 

deepmd/infer/deep_polar.py

@@ -94,7 +94,7 @@ def eval(
         coords: np.array,
         cells: np.array,
         atom_types: List[int],
-        atomic: bool = True,
+        atomic: bool = False,
         fparam: Optional[np.array] = None,
         aparam: Optional[np.array] = None,
         efield: Optional[np.array] = None,
@@ -114,7 +114,7 @@ def eval(
             The atom types
             The list should contain natoms ints
         atomic
-            Calculate the atomic energy and virial
+            Not used in this model
         fparam
             Not used in this model
         aparam

deepmd/infer/deep_tensor.py

@@ -1,5 +1,5 @@
 import os
-from typing import List, Optional, TYPE_CHECKING
+from typing import List, Optional, TYPE_CHECKING, Tuple
 
 import numpy as np
 from deepmd.common import make_default_mesh
@@ -53,10 +53,34 @@ def __init__(
             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)
+            self.tensors.update(optional_tensors)
+            self._support_gfv = True
+        except KeyError:
+            self._support_gfv = False
+
         # start a tf session associated to the graph
         self.sess = tf.Session(graph=self.graph, config=default_tf_session_config)
         self._run_default_sess()
@@ -117,7 +141,8 @@ def eval(
             The atom types
             The list should contain natoms ints
         atomic
-            Calculate the atomic energy and virial
+            If True (default), return the atomic tensor
+            Otherwise return the global tensor
         fparam
             Not used in this model
         aparam
@@ -133,15 +158,16 @@ def eval(
                 else of size nframes x natoms x output_dim
         """
         # standarize the shape of inputs
-        coords = np.array(coords)
-        cells = np.array(cells)
-        atom_types = np.array(atom_types, dtype = int)
-
-        # reshape the inputs 
-        cells = np.reshape(cells, [-1, 9])
-        nframes = cells.shape[0]
-        coords = np.reshape(coords, [nframes, -1])
-        natoms = coords.shape[1] // 3
+        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())
@@ -151,14 +177,24 @@ def eval(
         assert(natoms_vec[0] == natoms)
 
         # evaluate
-        tensor = []
         feed_dict_test = {}
         feed_dict_test[self.t_natoms] = natoms_vec
         feed_dict_test[self.t_type  ] = np.tile(atom_types, [nframes,1]).reshape([-1])
-        t_out = [self.t_tensor]
         feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
         feed_dict_test[self.t_box  ] = np.reshape(cells , [-1])
-        feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
+        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]
 
@@ -171,3 +207,122 @@ def eval(
             tensor = np.reshape(tensor, [nframes, self.output_dim])
 
         return tensor
+
+    def eval_full(
+        self,
+        coords: np.array,
+        cells: np.array,
+        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

deepmd/loss/tensor.py

@@ -56,7 +56,7 @@ def build (self,
                suffix):        
         polar_hat = label_dict[self.label_name]
         atomic_polar_hat = label_dict["atomic_" + self.label_name]
-        polar = model_dict[self.tensor_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]

deepmd/model/tensor.py

@@ -112,6 +112,9 @@ def build (self,
                                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()
+
         dout \
             = self.descrpt.build(coord_,
                                  atype_,
@@ -123,16 +126,49 @@ def build (self,
                                  reuse = reuse)
         dout = tf.identity(dout, name='o_descriptor')
         rot_mat = self.descrpt.get_rot_mat()
-        rot_mat = tf.identity(rot_mat, name = 'o_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)
-        output = tf.identity(output, name = 'o_' + self.model_type)
-
-        return {self.model_type: output}
+        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
 
 
 class WFCModel(TensorModel):