feat: calculate the real error in dp model-devi (#2757)

This PR adds a new argument `dp model-devi --real_error`. With that, the
real error given in the DP-GEN paper is calculated.

![Screenshot from 2023-08-23
15-43-57](https://github.com/deepmodeling/deepmd-kit/assets/9496702/331dd7fc-f1d4-4b37-be9f-80c78d027602)

One can use this feature for `dpgen simplify` if the labeled data
exists.

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

Author: njzjz

deepmd/infer/model_devi.py

@@ -21,13 +21,18 @@
 )
 
 
-def calc_model_devi_f(fs: np.ndarray) -> Tuple[np.ndarray]:
+def calc_model_devi_f(
+    fs: np.ndarray, real_f: Optional[np.ndarray] = None
+) -> Tuple[np.ndarray]:
     """Calculate model deviation of force.
 
     Parameters
     ----------
     fs : numpy.ndarray
         size of `n_models x n_frames x n_atoms x 3`
+    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.
 
     Returns
     -------
@@ -38,14 +43,21 @@ def calc_model_devi_f(fs: np.ndarray) -> Tuple[np.ndarray]:
     avg_devi_f : numpy.ndarray
         average deviation of force in all atoms
     """
-    fs_devi = np.linalg.norm(np.std(fs, axis=0), axis=-1)
+    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
+        )
     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) -> np.ndarray:
+def calc_model_devi_e(
+    es: np.ndarray, real_e: Optional[np.ndarray] = None
+) -> np.ndarray:
     """Calculate model deviation of total energy per atom.
 
     Here we don't use the atomic energy, as the decomposition
@@ -56,24 +68,35 @@ def calc_model_devi_e(es: np.ndarray) -> np.ndarray:
     ----------
     es : numpy.ndarray
         size of `n_models x n_frames x 1
+    real_e : numpy.ndarray
+        real energy, size of `n_frames x 1`. If given,
+        the RMS real error is calculated instead.
 
     Returns
     -------
     max_devi_e : numpy.ndarray
         maximum deviation of energy
     """
-    es_devi = np.std(es, axis=0)
+    if real_e is None:
+        es_devi = np.std(es, axis=0)
+    else:
+        es_devi = np.sqrt(np.mean(np.square(es - real_e), axis=0))
     es_devi = np.squeeze(es_devi, axis=-1)
     return es_devi
 
 
-def calc_model_devi_v(vs: np.ndarray) -> Tuple[np.ndarray]:
+def calc_model_devi_v(
+    vs: np.ndarray, real_v: Optional[np.ndarray] = None
+) -> Tuple[np.ndarray]:
     """Calculate model deviation of virial.
 
     Parameters
     ----------
     vs : numpy.ndarray
         size of `n_models x n_frames x 9`
+    real_v : numpy.ndarray
+        real virial, size of `n_frames x 9`. If given,
+        the RMS real error is calculated instead.
 
     Returns
     -------
@@ -84,7 +107,10 @@ def calc_model_devi_v(vs: np.ndarray) -> Tuple[np.ndarray]:
     avg_devi_v : numpy.ndarray
         average deviation of virial in 9 elements
     """
-    vs_devi = np.std(vs, axis=0)
+    if real_v is None:
+        vs_devi = np.std(vs, axis=0)
+    else:
+        vs_devi = np.sqrt(np.mean(np.square(vs - real_v), 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
@@ -148,6 +174,7 @@ def calc_model_devi(
     mixed_type=False,
     fparam: Optional[np.ndarray] = None,
     aparam: Optional[np.ndarray] = None,
+    real_data: Optional[dict] = None,
 ):
     """Python interface to calculate model deviation.
 
@@ -171,6 +198,8 @@ def calc_model_devi(
         frame specific parameters
     aparam : numpy.ndarray
         atomic specific parameters
+    real_data : dict, optional
+        real data to calculate RMS real error
 
     Returns
     -------
@@ -211,17 +240,29 @@ def calc_model_devi(
     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.append(calc_model_devi_e(energies))
+    if real_data is None:
+        devi += list(calc_model_devi_v(virials))
+        devi += list(calc_model_devi_f(forces))
+        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.append(calc_model_devi_e(energies, real_data["energy"]))
     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
+    *,
+    models: list,
+    system: str,
+    set_prefix: str,
+    output: str,
+    frequency: int,
+    real_error: bool = False,
+    **kwargs,
 ):
     """Make model deviation calculation.
 
@@ -239,6 +280,8 @@ def make_model_devi(
         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.
+    real_error : bool, default: False
+        If True, calculate the RMS real error instead of model deviation.
     **kwargs
         Arbitrary keyword arguments.
     """
@@ -279,6 +322,29 @@ def make_model_devi(
                 must=True,
                 high_prec=False,
             )
+        if real_error:
+            dp_data.add(
+                "energy",
+                1,
+                atomic=False,
+                must=False,
+                high_prec=True,
+            )
+            dp_data.add(
+                "force",
+                3,
+                atomic=True,
+                must=False,
+                high_prec=False,
+            )
+            dp_data.add(
+                "virial",
+                9,
+                atomic=False,
+                must=False,
+                high_prec=False,
+            )
+
         mixed_type = dp_data.mixed_type
 
         data_sets = [dp_data._load_set(set_name) for set_name in dp_data.dirs]
@@ -301,6 +367,15 @@ def make_model_devi(
                 aparam = data["aparam"]
             else:
                 aparam = None
+            if real_error:
+                natoms = atype.shape[-1]
+                real_data = {
+                    "energy": data["energy"] / natoms,
+                    "force": data["force"].reshape([-1, natoms, 3]),
+                    "virial": data["virial"] / natoms,
+                }
+            else:
+                real_data = None
             devi = calc_model_devi(
                 coord,
                 box,
@@ -309,6 +384,7 @@ def make_model_devi(
                 mixed_type=mixed_type,
                 fparam=fparam,
                 aparam=aparam,
+                real_data=real_data,
             )
             nframes_tot += coord.shape[0]
             devis.append(devi)

deepmd_cli/main.py

@@ -438,6 +438,12 @@ def main_parser() -> argparse.ArgumentParser:
         type=int,
         help="The trajectory frequency of the system",
     )
+    parser_model_devi.add_argument(
+        "--real_error",
+        action="store_true",
+        default=False,
+        help="Calculate the RMS real error of the model. The real data should be given in the systems.",
+    )
 
     # * convert models
     parser_transform = subparsers.add_parser(

source/tests/test_model_devi.py

@@ -86,6 +86,33 @@ def test_make_model_devi(self):
         x = np.loadtxt(self.output)
         np.testing.assert_allclose(x, self.expect, 6)
 
+    def test_make_model_devi_real_erorr(self):
+        make_model_devi(
+            models=self.graph_dirs,
+            system=self.data_dir,
+            set_prefix="set",
+            output=self.output,
+            frequency=self.freq,
+            real_error=True,
+        )
+        x = np.loadtxt(self.output)
+        np.testing.assert_allclose(
+            x,
+            np.array(
+                [
+                    0.000000e00,
+                    6.709021e-01,
+                    1.634359e-03,
+                    3.219720e-01,
+                    2.018684e00,
+                    1.829748e00,
+                    1.956474e00,
+                    1.550898e02,
+                ]
+            ),
+            6,
+        )
+
     def tearDown(self):
         for pb in self.graph_dirs:
             os.remove(pb)