support atomic/relative model deviation in CLI

deepmd/infer/model_devi.py

@@ -2,6 +2,7 @@
 from typing import (
     Optional,
     Tuple,
+    overload,
 )
 
 import numpy as np
@@ -20,10 +21,39 @@
     DeepPot,
 )
 
+try:
+    from typing import Literal  # python >=3.8
+except ImportError:
+    from typing_extensions import Literal  # type: ignore
+
+
+@overload
+def calc_model_devi_f(
+    fs: np.ndarray,
+    real_f: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
+    atomic: Literal[False] = False,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    ...
+
+
+@overload
+def calc_model_devi_f(
+    fs: np.ndarray,
+    real_f: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
+    *,
+    atomic: Literal[True],
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    ...
+
 
 def calc_model_devi_f(
-    fs: np.ndarray, real_f: Optional[np.ndarray] = None
-) -> Tuple[np.ndarray]:
+    fs: np.ndarray,
+    real_f: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
+    atomic: bool = False,
+) -> Tuple[np.ndarray, ...]:
     """Calculate model deviation of force.
 
     Parameters
@@ -33,6 +63,12 @@ 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.
+    atomic : bool, default: False
+        Whether return deviation of force in all atoms
 
     Returns
     -------
@@ -42,16 +78,30 @@ 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, returned if atomic=True
     """
     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)
+    if atomic:
+        return max_devi_f, min_devi_f, avg_devi_f, fs_devi
     return max_devi_f, min_devi_f, avg_devi_f
 
 
@@ -86,8 +136,10 @@ def calc_model_devi_e(
 
 
 def calc_model_devi_v(
-    vs: np.ndarray, real_v: Optional[np.ndarray] = None
-) -> Tuple[np.ndarray]:
+    vs: np.ndarray,
+    real_v: Optional[np.ndarray] = None,
+    relative: Optional[float] = None,
+) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Calculate model deviation of virial.
 
     Parameters
@@ -97,6 +149,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 +167,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 +196,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 +211,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 +250,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 +278,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 +329,26 @@ 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, atomic=atomic))
+        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, atomic=atomic
+            )
+        )
+        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 +360,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 +383,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 +416,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 +498,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

pyproject.toml

@@ -43,7 +43,7 @@ dependencies = [
     'pyyaml',
     'dargs >= 0.3.5',
     'python-hostlist >= 1.21',
-    'typing_extensions; python_version < "3.7"',
+    'typing_extensions; python_version < "3.8"',
     'importlib_metadata>=1.4; python_version < "3.8"',
     'h5py',
     'wcmatch',