support atomic/relative model deviation in CLI

Fix deepmodeling#2017.

Signed-off-by: Jinzhe Zeng <[email protected]>

deepmd/infer/model_devi.py

@@ -22,7 +22,9 @@
 
 
 def calc_model_devi_f(
-    fs: np.ndarray, real_f: Optional[np.ndarray] = None
+    fs: np.ndarray,
+    real_f: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
 ) -> Tuple[np.ndarray]:
     """Calculate model deviation of force.
 
@@ -33,6 +35,10 @@ def calc_model_devi_f(
     real_f : numpy.ndarray or None
         real force, size of `n_frames x n_atoms x 3`. If given,
         the RMS real error is calculated instead.
+    relative : float, default: None
+        If given, calculate the relative model deviation of force. The
+        value is the level parameter for computing the relative model
+        deviation of the force.
 
     Returns
     -------
@@ -42,17 +48,29 @@ def calc_model_devi_f(
         minimum deviation of force in all atoms
     avg_devi_f : numpy.ndarray
         average deviation of force in all atoms
+    fs_devi : numpy.ndarray
+        deviation of force in all atoms
     """
     if real_f is None:
         fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1)
     else:
         fs_devi = np.linalg.norm(
             np.sqrt(np.mean(np.square(fs - real_f), axis=0)), axis=-1
         )
+    if relative is not None:
+        if real_f is None:
+            # if real force is not given, the magnitude is calculated from mean value of four models
+            # See DeepPotModelDevi::compute_relative_std_f
+            # See also Eq. 71 in DeePMD-kit v2 paepr
+            magnitude = np.linalg.norm(np.mean(fs, axis=0), axis=-1)
+        else:
+            # otherwise, the magnitude is calculated from the real force
+            magnitude = np.linalg.norm(real_f, axis=-1)
+        fs_devi /= magnitude + relative
     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
+    return max_devi_f, min_devi_f, avg_devi_f, fs_devi
 
 
 def calc_model_devi_e(
@@ -86,7 +104,9 @@ def calc_model_devi_e(
 
 
 def calc_model_devi_v(
-    vs: np.ndarray, real_v: Optional[np.ndarray] = None
+    vs: np.ndarray,
+    real_v: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
 ) -> Tuple[np.ndarray]:
     """Calculate model deviation of virial.
 
@@ -97,6 +117,10 @@ def calc_model_devi_v(
     real_v : numpy.ndarray
         real virial, size of `n_frames x 9`. If given,
         the RMS real error is calculated instead.
+    relative : float, default: None
+        If given, calculate the relative model deviation of virial. The
+        value is the level parameter for computing the relative model
+        deviation of the virial.
 
     Returns
     -------
@@ -111,13 +135,25 @@ def calc_model_devi_v(
         vs_devi = np.std(vs, axis=0)
     else:
         vs_devi = np.sqrt(np.mean(np.square(vs - real_v), axis=0))
+    if relative is not None:
+        if real_v is None:
+            # if real virial is not given, the magnitude is calculated from mean value of four models
+            # See DeepPotModelDevi::compute_relative_std_v
+            # See also Eq. 72 in DeePMD-kit v2 paepr
+            magnitude = np.linalg.norm(np.mean(vs, axis=0), axis=-1)
+        else:
+            # otherwise, the magnitude is calculated from the real virial
+            magnitude = np.linalg.norm(real_v, axis=-1)
+        vs_devi /= magnitude + relative
     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 = ""):
+def write_model_devi_out(
+    devi: np.ndarray, fname: str, header: str = "", atomic: bool = False
+):
     """Write output of model deviation.
 
     Parameters
@@ -128,8 +164,13 @@ def write_model_devi_out(devi: np.ndarray, fname: str, header: str = ""):
         the file name to dump
     header : str, default=""
         the header to dump
+    atomic : bool, default: False
+        whether atomic model deviation is printed
     """
-    assert devi.shape[1] == 8
+    if not atomic:
+        assert devi.shape[1] == 8
+    else:
+        assert devi.shape[1] > 8
     header = "%s\n%10s" % (header, "step")
     for item in "vf":
         header += "%19s%19s%19s" % (
@@ -138,11 +179,13 @@ def write_model_devi_out(devi: np.ndarray, fname: str, header: str = ""):
             f"avg_devi_{item}",
         )
     header += "%19s" % "devi_e"
+    if atomic:
+        header += "%19s" % "atm_devi_f(N)"
     with open(fname, "ab") as fp:
         np.savetxt(
             fp,
             devi,
-            fmt=["%12d"] + ["%19.6e" for _ in range(7)],
+            fmt=["%12d"] + ["%19.6e" for _ in range(devi.shape[1] - 1)],
             delimiter="",
             header=header,
         )
@@ -175,6 +218,9 @@ def calc_model_devi(
     fparam: Optional[np.ndarray] = None,
     aparam: Optional[np.ndarray] = None,
     real_data: Optional[dict] = None,
+    atomic: bool = False,
+    relative: Optional[float] = None,
+    relative_v: Optional[float] = None,
 ):
     """Python interface to calculate model deviation.
 
@@ -200,6 +246,16 @@ def calc_model_devi(
         atomic specific parameters
     real_data : dict, optional
         real data to calculate RMS real error
+    atomic : bool, default: False
+        If True, calculate the force model deviation of each atom.
+    relative : float, default: None
+        If given, calculate the relative model deviation of force. The
+        value is the level parameter for computing the relative model
+        deviation of the force.
+    relative_v : float, default: None
+        If given, calculate the relative model deviation of virial. The
+        value is the level parameter for computing the relative model
+        deviation of the virial.
 
     Returns
     -------
@@ -241,16 +297,22 @@ def calc_model_devi(
 
     devi = [np.arange(coord.shape[0]) * frequency]
     if real_data is None:
-        devi += list(calc_model_devi_v(virials))
-        devi += list(calc_model_devi_f(forces))
+        devi += list(calc_model_devi_v(virials, relative=relative_v))
+        devi_f = list(calc_model_devi_f(forces, relative=relative))
+        devi += devi_f[:3]
         devi.append(calc_model_devi_e(energies))
     else:
-        devi += list(calc_model_devi_v(virials, real_data["virial"]))
-        devi += list(calc_model_devi_f(forces, real_data["force"]))
+        devi += list(
+            calc_model_devi_v(virials, real_data["virial"], relative=relative_v)
+        )
+        devi_f = list(calc_model_devi_f(forces, real_data["force"], relative=relative))
+        devi += devi_f[:3]
         devi.append(calc_model_devi_e(energies, real_data["energy"]))
     devi = np.vstack(devi).T
+    if atomic:
+        devi = np.concatenate([devi, devi_f[3]], axis=1)
     if fname:
-        write_model_devi_out(devi, fname)
+        write_model_devi_out(devi, fname, atomic=atomic)
     return devi
 
 
@@ -262,6 +324,9 @@ def make_model_devi(
     output: str,
     frequency: int,
     real_error: bool = False,
+    atomic: bool = False,
+    relative: Optional[float] = None,
+    relative_v: Optional[float] = None,
     **kwargs,
 ):
     """Make model deviation calculation.
@@ -282,6 +347,16 @@ def make_model_devi(
         This paramter is used to determine the index in the output file.
     real_error : bool, default: False
         If True, calculate the RMS real error instead of model deviation.
+    atomic : bool, default: False
+        If True, calculate the force model deviation of each atom.
+    relative : float, default: None
+        If given, calculate the relative model deviation of force. The
+        value is the level parameter for computing the relative model
+        deviation of the force.
+    relative_v : float, default: None
+        If given, calculate the relative model deviation of virial. The
+        value is the level parameter for computing the relative model
+        deviation of the virial.
     **kwargs
         Arbitrary keyword arguments.
     """
@@ -305,7 +380,9 @@ def make_model_devi(
 
     for system in all_sys:
         # create data-system
-        dp_data = DeepmdData(system, set_prefix, shuffle_test=False, type_map=tmap)
+        dp_data = DeepmdData(
+            system, set_prefix, shuffle_test=False, type_map=tmap, sort_atoms=False
+        )
         if first_dp.get_dim_fparam() > 0:
             dp_data.add(
                 "fparam",
@@ -385,11 +462,14 @@ def make_model_devi(
                 fparam=fparam,
                 aparam=aparam,
                 real_data=real_data,
+                atomic=atomic,
+                relative=relative,
+                relative_v=relative_v,
             )
             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)
+        write_model_devi_out(devis, output, header=system, atomic=atomic)
         devis_coll.append(devis)
     return devis_coll

deepmd_cli/main.py

@@ -448,6 +448,22 @@ def main_parser() -> argparse.ArgumentParser:
         default=False,
         help="Calculate the RMS real error of the model. The real data should be given in the systems.",
     )
+    parser_model_devi.add_argument(
+        "--atomic",
+        action="store_true",
+        default=False,
+        help="Print the force model deviation of each atom.",
+    )
+    parser_model_devi.add_argument(
+        "--relative",
+        type=float,
+        help="Calculate the relative model deviation of force. The level parameter for computing the relative model deviation of the force should be given.",
+    )
+    parser_model_devi.add_argument(
+        "--relative_v",
+        type=float,
+        help="Calculate the relative model deviation of virial. The level parameter for computing the relative model deviation of the virial should be given.",
+    )
 
     # * convert models
     parser_transform = subparsers.add_parser(

doc/test/model-deviation.md

@@ -36,3 +36,14 @@ optional arguments:
 ```
 
 For more details concerning the definition of model deviation and its application, please refer to [Yuzhi Zhang, Haidi Wang, Weijie Chen, Jinzhe Zeng, Linfeng Zhang, Han Wang, and Weinan E, DP-GEN: A concurrent learning platform for the generation of reliable deep learning based potential energy models, Computer Physics Communications, 2020, 253, 107206.](https://doi.org/10.1016/j.cpc.2020.107206)
+
+## Relative model deviation
+
+By default, the model deviation is output in absolute value. If the argument `--relative` is passed, then the relative model deviation of the force will be output, including values output by the argument `--atomic`. The relative model deviation of the force on atom $i$ is defined by
+
+$$E_{f_i}=\frac{\left|D_{f_i}\right|}{\left|f_i\right|+l}$$
+
+where $D_{f_i}$ is the absolute model deviation of the force on atom $i$, $f_i$ is the norm of the force and $l$ is provided as the parameter of the keyword `relative`.
+If the argument `--relative_v` is set, then the relative model deviation of the virial will be output instead of the absolute value, with the same definition of that of the force:
+
+$$E_{v_i}=\frac{\left|D_{v_i}\right|}{\left|v_i\right|+l}$$

doc/third-party/lammps-command.md

@@ -40,7 +40,7 @@ and the model deviation will be computed among all models every `out_freq` times
     <i>fparam_from_compute</i> value = id
         id = compute id used to update the frame parameter.
     <i>atomic</i> = no value is required.
-        If this keyword is set, the model deviation of each atom will be output.
+        If this keyword is set, the force model deviation of each atom will be output.
     <i>relative</i> value = level
         level = The level parameter for computing the relative model deviation of the force
     <i>relative_v</i> value = level

source/tests/test_model_devi.py

@@ -113,6 +113,95 @@ def test_make_model_devi_real_erorr(self):
             6,
         )
 
+    def test_make_model_devi_atomic_relative(self):
+        _, expected_f, expected_v = self.graphs[0].eval(
+            self.coord[0], self.box[0], self.atype
+        )
+        _, expected_f2, expected_v2 = self.graphs[1].eval(
+            self.coord[0], self.box[0], self.atype
+        )
+        expected_f = expected_f.reshape((-1, 3))
+        expected_f2 = expected_f2.reshape((-1, 3))
+        expected_v = expected_v.reshape((-1, 3))
+        expected_v2 = expected_v2.reshape((-1, 3))
+        relative = 1.0
+        make_model_devi(
+            models=self.graph_dirs,
+            system=self.data_dir,
+            set_prefix="set",
+            output=self.output,
+            frequency=self.freq,
+            atomic=True,
+            relative=relative,
+        )
+        md = np.loadtxt(self.output)
+        # copy from lammps test
+        norm = np.linalg.norm(np.mean([expected_f, expected_f2], axis=0), axis=1)
+        expected_md_f = np.linalg.norm(
+            np.std([expected_f, expected_f2], axis=0), axis=1
+        )
+        expected_md_f /= norm + relative
+        np.testing.assert_allclose(md[8:], expected_md_f, 6)
+        np.testing.assert_allclose(md[7], self.expect[7], 6)
+        np.testing.assert_allclose(md[4], np.max(expected_md_f), 6)
+        np.testing.assert_allclose(md[5], np.min(expected_md_f), 6)
+        np.testing.assert_allclose(md[6], np.mean(expected_md_f), 6)
+        expected_md_v = (
+            np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0)
+            / 6
+        )
+        np.testing.assert_allclose(md[1], np.max(expected_md_v), 6)
+        np.testing.assert_allclose(md[2], np.min(expected_md_v), 6)
+        np.testing.assert_allclose(md[3], np.sqrt(np.mean(np.square(expected_md_v))), 6)
+
+    def test_make_model_devi_atomic_relative_v(self):
+        _, expected_f, expected_v = self.graphs[0].eval(
+            self.coord[0], self.box[0], self.atype
+        )
+        _, expected_f2, expected_v2 = self.graphs[1].eval(
+            self.coord[0], self.box[0], self.atype
+        )
+        expected_f = expected_f.reshape((-1, 3))
+        expected_f2 = expected_f2.reshape((-1, 3))
+        expected_v = expected_v.reshape((-1, 3))
+        expected_v2 = expected_v2.reshape((-1, 3))
+        relative = 1.0
+        make_model_devi(
+            models=self.graph_dirs,
+            system=self.data_dir,
+            set_prefix="set",
+            output=self.output,
+            frequency=self.freq,
+            atomic=True,
+            relative_v=relative,
+        )
+        md = np.loadtxt(self.output)
+        # copy from lammps test
+        expected_md_f = np.linalg.norm(
+            np.std([expected_f, expected_f2], axis=0), axis=1
+        )
+        np.testing.assert_allclose(md[8:], expected_md_f, 6)
+        np.testing.assert_allclose(md[7], self.expect[7], 6)
+        np.testing.assert_allclose(md[4], np.max(expected_md_f), 6)
+        np.testing.assert_allclose(md[5], np.min(expected_md_f), 6)
+        np.testing.assert_allclose(md[6], np.mean(expected_md_f), 6)
+        expected_md_v = (
+            np.std([np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0)
+            / 6
+        )
+        norm = (
+            np.abs(
+                np.mean(
+                    [np.sum(expected_v, axis=0), np.sum(expected_v2, axis=0)], axis=0
+                )
+            )
+            / 6
+        )
+        expected_md_v /= norm + relative
+        np.testing.assert_allclose(md[1], np.max(expected_md_v), 6)
+        np.testing.assert_allclose(md[2], np.min(expected_md_v), 6)
+        np.testing.assert_allclose(md[3], np.sqrt(np.mean(np.square(expected_md_v))), 6)
+
     def tearDown(self):
         for pb in self.graph_dirs:
             os.remove(pb)