Support relative force model deviation by normalizing the RMS force magnitude (#496)

* support relative force model deviation normalized by avg f of MD trajectories

* update readme

Co-authored-by: Han Wang <wang_han@iapcm.ac.cn>

Author: amcadmus

README.md

@@ -556,6 +556,7 @@ The bold notation of key (such aas **type_map**) means that it's a necessary key
 | model_devi_numb_candi_v  | Int | 0  | See `model_devi_adapt_trust_lo`.|
 | model_devi_perc_candi_f  | Float | 0.0 | See `model_devi_adapt_trust_lo`.|
 | model_devi_perc_candi_v  | Float | 0.0 | See `model_devi_adapt_trust_lo`.|
+| model_devi_f_avg_relative | Boolean | False | Normalized the force model deviations by the RMS force magnitude along the trajectory. This key should not be used with `use_relative`. |
 | **model_devi_clean_traj**  | Boolean | true                                                         | Deciding whether to clean traj folders in MD since they are too large. |
 | **model_devi_nopbc**  | Boolean | False                                                         | Assume open boundary condition in MD simulations. |
 | model_devi_activation_func | List of list of string | [["tanh","tanh"],["tanh","gelu"],["gelu","tanh"],["gelu","gelu"]]	| Set activation functions for models, length of the List should be the same as `numb_models`, and two elements in the list of string respectively assign activation functions to the embedding and fitting nets within each model. *Backward compatibility*: the orginal "List of String" format is still supported, where embedding and fitting nets of one model use the same activation function, and the length of the List should be the same as `numb_models`|

dpgen/generator/lib/lammps.py

@@ -89,7 +89,7 @@ def make_lammps_input(ensemble,
     ret+= "\n"
     ret+= "thermo_style    custom step temp pe ke etotal press vol lx ly lz xy xz yz\n"
     ret+= "thermo          ${THERMO_FREQ}\n"
-    ret+= "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z\n"
+    ret+= "dump            1 all custom ${DUMP_FREQ} traj/*.lammpstrj id type x y z fx fy fz\n"
     ret+= "restart         10000 dpgen.restart\n"
     ret+= "\n"
     if pka_e is None :
@@ -137,5 +137,37 @@ def make_lammps_input(ensemble,
 # cvt_lammps_conf('POSCAR', 'tmp.lmp')
 
 
-    
- 
+def get_dumped_forces(
+        file_name):
+    with open(file_name) as fp:        
+        lines = fp.read().split('\n')
+    natoms = None
+    for idx,ii in enumerate(lines):
+        if 'ITEM: NUMBER OF ATOMS' in ii:
+            natoms = int(lines[idx+1])
+            break
+    if natoms is None:
+        raise RuntimeError('wrong dump file format, cannot find number of atoms', file_name)
+    idfx = None
+    for idx,ii in enumerate(lines):
+        if 'ITEM: ATOMS' in ii:
+            keys = ii
+            keys = keys.replace('ITEM: ATOMS', '')
+            keys = keys.split()
+            idfx = keys.index('fx')
+            idfy = keys.index('fy')
+            idfz = keys.index('fz')
+            break
+    if idfx is None:
+        raise RuntimeError('wrong dump file format, cannot find dump keys', file_name)
+    ret = []
+    for ii in range(idx+1, idx+natoms+1):
+        words = lines[ii].split()
+        ret.append([ float(words[ii]) for ii in [idfx, idfy, idfz] ])
+    ret = np.array(ret)
+    return ret
+
+
+if __name__ == '__main__':
+    ret = get_dumped_forces('40.lammpstrj')
+    print(ret)

dpgen/generator/run.py

@@ -40,7 +40,7 @@
 from dpgen.generator.lib.utils import record_iter
 from dpgen.generator.lib.utils import log_task
 from dpgen.generator.lib.utils import symlink_user_forward_files
-from dpgen.generator.lib.lammps import make_lammps_input
+from dpgen.generator.lib.lammps import make_lammps_input, get_dumped_forces
 from dpgen.generator.lib.vasp import write_incar_dict
 from dpgen.generator.lib.vasp import make_vasp_incar_user_dict
 from dpgen.generator.lib.vasp import incar_upper
@@ -1323,6 +1323,24 @@ def check_bad_box(conf_name,
     return is_bad
 
 
+def _read_model_devi_file(
+        task_path : str,
+        model_devi_f_avg_relative : bool = False
+):
+    model_devi = np.loadtxt(os.path.join(task_path, 'model_devi.out'))
+    if model_devi_f_avg_relative :
+        trajs = glob.glob(os.path.join(task_path, 'traj', '*.lammpstrj'))
+        all_f = []
+        for ii in trajs:
+            all_f.append(get_dumped_forces(ii))
+        all_f = np.array(all_f)
+        all_f = all_f.reshape([-1,3])
+        avg_f = np.sqrt(np.average(np.sum(np.square(all_f), axis = 1)))
+        model_devi[:,4:7] = model_devi[:,4:7] / avg_f
+        np.savetxt(os.path.join(task_path, 'model_devi_avgf.out'), model_devi, fmt='%16.6e')
+    return model_devi
+
+
 def _select_by_model_devi_standard(
         modd_system_task: List[str],
         f_trust_lo : float,
@@ -1331,6 +1349,7 @@ def _select_by_model_devi_standard(
         v_trust_hi : float,
         cluster_cutoff : float, 
         model_devi_skip : int = 0,
+        model_devi_f_avg_relative : bool = False,
         detailed_report_make_fp : bool = True,
 ):
     fp_candidate = []
@@ -1345,7 +1364,7 @@ def _select_by_model_devi_standard(
     for tt in modd_system_task :
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            all_conf = np.loadtxt(os.path.join(tt, 'model_devi.out'))
+            all_conf = _read_model_devi_file(tt, model_devi_f_avg_relative)
             for ii in range(all_conf.shape[0]) :
                 if all_conf[ii][0] < model_devi_skip :
                     continue
@@ -1393,7 +1412,8 @@ def _select_by_model_devi_adaptive_trust_low(
         v_trust_hi : float,
         numb_candi_v : int,
         perc_candi_v : float,
-        model_devi_skip : int = 0
+        model_devi_skip : int = 0,
+        model_devi_f_avg_relative : bool = False,
 ):
     """
     modd_system_task    model deviation tasks belonging to one system
@@ -1424,6 +1444,7 @@ def _select_by_model_devi_adaptive_trust_low(
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
             model_devi = np.loadtxt(os.path.join(tt, 'model_devi.out'))
+            model_devi = _read_model_devi_file(tt, model_devi_f_avg_relative)
             for ii in range(model_devi.shape[0]) :
                 if model_devi[ii][0] < model_devi_skip :
                     continue
@@ -1512,6 +1533,7 @@ def _make_fp_vasp_inner (modd_path,
     fp_tasks = []
     cluster_cutoff = jdata['cluster_cutoff'] if jdata.get('use_clusters', False) else None
     model_devi_adapt_trust_lo = jdata.get('model_devi_adapt_trust_lo', False)
+    model_devi_f_avg_relative = jdata.get('model_devi_f_avg_relative', False)
     # skip save *.out if detailed_report_make_fp is False, default is True
     detailed_report_make_fp = jdata.get("detailed_report_make_fp", True)
     # skip bad box criteria
@@ -1532,7 +1554,9 @@ def _make_fp_vasp_inner (modd_path,
                     v_trust_lo, v_trust_hi,
                     cluster_cutoff, 
                     model_devi_skip,
-                    detailed_report_make_fp = detailed_report_make_fp)
+                    model_devi_f_avg_relative = model_devi_f_avg_relative,
+                    detailed_report_make_fp = detailed_report_make_fp,
+                )
         else:
             numb_candi_f = jdata.get('model_devi_numb_candi_f', 10)
             numb_candi_v = jdata.get('model_devi_numb_candi_v', 0)
@@ -1543,7 +1567,9 @@ def _make_fp_vasp_inner (modd_path,
                     modd_system_task,
                     f_trust_hi, numb_candi_f, perc_candi_f,
                     v_trust_hi, numb_candi_v, perc_candi_v,
-                    model_devi_skip = model_devi_skip)
+                    model_devi_skip = model_devi_skip,
+                    model_devi_f_avg_relative = model_devi_f_avg_relative,
+                )
             dlog.info("system {0:s} {1:9s} : f_trust_lo {2:6.3f}   v_trust_lo {3:6.3f}".format(ss, 'adapted', f_trust_lo_ad, v_trust_lo_ad))
 
         # print a report

tests/generator/context.py

@@ -4,6 +4,7 @@
 from dpgen.generator.run import *
 from dpgen.generator.lib.gaussian import detect_multiplicity
 from dpgen.generator.lib.ele_temp import NBandsEsti
+from dpgen.generator.lib.lammps import get_dumped_forces
 
 param_file = 'param-mg-vasp.json'
 param_file_v1 = 'param-mg-vasp-v1.json'

tests/generator/test_lammps.py

@@ -0,0 +1,38 @@
+import os,sys,json,glob,shutil,textwrap
+import dpdata
+import numpy as np
+import unittest
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
+__package__ = 'generator'
+from .context import get_dumped_forces
+
+class TestGetDumpForce(unittest.TestCase):
+    def setUp(self):
+        file_content = textwrap.dedent("""\
+ITEM: TIMESTEP
+40
+ITEM: NUMBER OF ATOMS
+2
+ITEM: BOX BOUNDS xy xz yz pp pp pp
+-2.9180686220264818e-04 8.0855380329747089e+00 1.4011011277606830e-07
+-2.9198257591541018e-04 8.0855378881632269e+00 3.3202396460852749e-08
+-2.9180686326490957e-04 8.0855378891632768e+00 -1.7571268247505500e-07
+ITEM: ATOMS id type x y z fx fy fz
+1 1 2.09532 8.19528 2.00538 -0.00569269 -0.0200373 -0.0342394
+2 1 -0.0727384 4.01773 4.05582 -0.0297083 0.0817184 0.0722508
+""")
+        with open('tmp.dump', 'w') as fp:
+            fp.write(file_content)
+        self.expected_f = [ -0.00569269, -0.0200373, -0.0342394, -0.0297083, 0.0817184, 0.0722508]
+
+    def tearDown(self):
+        if os.path.isfile('tmp.dump'):
+            os.remove('tmp.dump')
+
+    def test_read_dump(self):
+        ff = get_dumped_forces('tmp.dump')
+        self.assertEqual(ff.shape, (2, 3))
+        ff = ff.reshape([-1])
+        for ii in range(6):
+            self.assertAlmostEqual(ff[ii], self.expected_f[ii])