Repository: ruoyuwang1995nya/pfd-kit

Issue:
validation test RMSE 
The calculated RMSE is smaller than that obtained by dptest

Hints from issue comments:
@ruoyuwang1995nya: Correct RMSE calculator has been implemented @fishssh 

Solution patch:
diff --git a/pfd/entrypoint/args.py b/pfd/entrypoint/args.py
index c36974a..aadcbea 100644
--- a/pfd/entrypoint/args.py
+++ b/pfd/entrypoint/args.py
@@ -70,6 +70,11 @@ def conf_args():
 def pert_gen():
     doc_atom_pert_distance = "Perturb distance for atoms, in Angstrom"
     doc_pert_num = "Number of perturbed structures"
+    doc_cell_pert = "The amount of lattice contraction or extension, relative to original lattice constant."
+    doc_replicate = (
+        "Generate supercell by lattice replication. Either an integer number or a list of three integers. "
+        "If an integer is given, the lattice is replicated uniformly in all three directions."
+    )
     return [
         Argument("conf_idx", [str, List[int]], optional=True, default="default"),
         Argument(
@@ -79,9 +84,13 @@ def pert_gen():
             default=0.0,
             doc=doc_atom_pert_distance,
         ),
-        Argument("cell_pert_fraction", float, optional=True, default=0.0),
+        Argument(
+            "cell_pert_fraction", float, optional=True, default=0.0, doc=doc_cell_pert
+        ),
         Argument("pert_num", int, optional=True, default=1, doc=doc_pert_num),
-        Argument("replicate", [int, List[int]], optional=True, default=1),
+        Argument(
+            "replicate", [int, List[int]], optional=True, default=1, doc=doc_replicate
+        ),
     ]
 
 
@@ -96,7 +105,7 @@ def normalize_pert_gen(data):
 
 def conf_gen_args():
     doc_init_conf = "The initial configurations for PFD workflow"
-    doc_pert_gen = "Generation of perturbed structures for PFD exploration"
+    doc_pert_gen = "Structure perturbation settings. A list of multiple perturbation settings can also be supplied if neccesarry."
     return [
         Argument(
             "init_confs",
diff --git a/pfd/exploration/inference/eval_model.py b/pfd/exploration/inference/eval_model.py
index 5ab9e75..74a04e2 100644
--- a/pfd/exploration/inference/eval_model.py
+++ b/pfd/exploration/inference/eval_model.py
@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from dataclasses import dataclass
+from math import sqrt
 from pathlib import Path
 from typing import Optional, Union
 import dpdata
@@ -62,14 +63,23 @@ def add_report(self, report: TestReport):
 
     def get_weighted_rmse_f(self):
         if len(self._reports) > 0:
-            return sum(res.numb_frame * res.rmse_f for res in self._reports) / sum(
-                res.numb_frame for res in self._reports
+            return np.sqrt(
+                sum(
+                    res.numb_frame * res.rmse_f**2 * 3 * res.atom_numb
+                    for res in self._reports
+                )
+                / sum(
+                    res.numb_frame * 3 * res.atom_numb
+                    for res in self._reports
+                    if res.atom_numb > 0
+                )
             )
 
     def get_weighted_rmse_e_atom(self):
         if len(self._reports) > 0:
-            return sum(res.numb_frame * res.rmse_e_atom for res in self._reports) / sum(
-                res.numb_frame for res in self._reports
+            return np.sqrt(
+                sum(res.numb_frame * res.rmse_e_atom**2 for res in self._reports)
+                / sum(res.numb_frame for res in self._reports)
             )
 
     def get_systems(self):


Test patch:
diff --git a/tests/test_test_reports.py b/tests/test_test_reports.py
new file mode 100644
index 0000000..a18f109
--- /dev/null
+++ b/tests/test_test_reports.py
@@ -0,0 +1,73 @@
+import sys
+
+sys.path.append("../")
+import unittest
+import copy
+from math import sqrt
+from pfd.op.stage import StageScheduler
+import dpdata
+from pathlib import Path
+import numpy as np
+from dflow.python import OPIO
+from pfd.exploration import explore_styles
+from pfd.exploration.scheduler import Scheduler
+from pfd.exploration.inference import TestReport, TestReports
+
+
+class TestTestReport(unittest.TestCase):
+    def setUp(self):
+        self.rep0 = TestReport(
+            name="sys0",
+            atom_numb=1,
+            numb_frame=1,
+            mae_e_atom=0.01,
+            mae_f=0.1,
+            rmse_e=0.02,
+            rmse_f=0.2,
+        )
+
+        self.rep1 = TestReport(
+            name="sys1",
+            atom_numb=2,
+            numb_frame=2,
+            mae_e_atom=0.02,
+            mae_f=0.2,
+            rmse_e=0.04,
+            rmse_f=0.4,
+        )
+
+        self.reports = TestReports()
+        self.reports.add_report(self.rep0)
+        self.reports.add_report(self.rep1)
+
+    def test_test_report(self):
+        rmse_f = sqrt(
+            (
+                3 * self.rep0.numb_frame * self.rep0.atom_numb * self.rep0.rmse_f**2
+                + 3 * self.rep1.numb_frame * self.rep1.atom_numb * self.rep1.rmse_f**2
+            )
+            / (
+                3 * self.rep0.numb_frame * self.rep0.atom_numb
+                + 3 * self.rep1.numb_frame * self.rep1.atom_numb
+            )
+        )
+        self.assertAlmostEqual(self.reports.get_weighted_rmse_f(), rmse_f)
+
+        rmse_e = sqrt(
+            (
+                self.rep0.numb_frame * self.rep0.atom_numb * self.rep0.rmse_e_atom**2
+                + self.rep1.numb_frame
+                * self.rep1.atom_numb
+                * self.rep1.rmse_e_atom**2
+            )
+            / (
+                3 * self.rep0.numb_frame * self.rep0.atom_numb
+                + 3 * self.rep1.numb_frame * self.rep1.atom_numb
+            )
+        )
+
+        self.assertAlmostEqual(self.reports.get_weighted_rmse_e_atom(), rmse_e)
+
+
+if __name__ == "__main__":
+    unittest.main()