fix dplr: support initializing from a restart file (deepmodeling#4084)

Thie PR aims at fixing deepmodeling#3679.

The `get_valid_pairs` will raise an error when
1) reading a restart file written from previous dplr jobs
2) reading a data file written from previous dplr jobs
3) executing 'run 0' for twice

The origin is that the check in `get_valid_pairs` is too strict for the
initialization process. The check requires that a pair of real and
virtual atoms found by iterating over the bondlist to be the local atoms
of the same processor. This requirement is automatically satisfied
during the running of MD because the positions of virtual atoms are set
equal to those of real atoms before rebuilding bondlists. However,
during initialization (`Verlet::setup(int flag)`), the bondlist is built
based on the initial configuration provided by the user. If the user
provides a configuration in which the virtual atoms are placed at the
Wannier centroids, it is possible that a bond formed by a pair of real
and virtual atoms are between a local and a ghost atom, which triggers
this error.

This PR loosens the requirement on the initial configuration provided by
the user. In the `setup_post_neighbor` function, the `get_valid_pairs`
function searches for bonded pairs but skips the aforementioned check.
It sets the coordinates of the virtual atoms to the coordinates of the
real atoms, and then rebuilds the bondlist. Afterwards, the check can be
passed.

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Introduced a new method for setting up pre-exchange operations,
enhancing simulation lifecycle management.
- **Changes**
- Renamed the `post_integrate` method to `pre_exchange` to better
reflect its role in the simulation process.
- Updated mask settings to improve clarity and context within the
simulation lifecycle.
- **Chores**
- Updated ownership context for Docker image in the build workflow
configuration.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

source/lmp/fix_dplr.cpp

@@ -236,7 +236,8 @@ int FixDPLR::setmask() {
   // THERMO_ENERGY removed in lammps/lammps#2560
   mask |= THERMO_ENERGY;
 #endif
-  mask |= POST_INTEGRATE;
+  mask |= PRE_EXCHANGE;
+  mask |= POST_NEIGHBOR;
   mask |= PRE_FORCE;
   mask |= POST_FORCE;
   mask |= MIN_PRE_EXCHANGE;
@@ -248,19 +249,9 @@ int FixDPLR::setmask() {
 /* ---------------------------------------------------------------------- */
 
 void FixDPLR::init() {
-  // double **xx = atom->x;
-  // double **vv = atom->v;
-  // int nlocal = atom->nlocal;
-  // for (int ii = 0; ii < nlocal; ++ii){
-  //   cout << xx[ii][0] << " "
-  // 	 << xx[ii][1] << " "
-  // 	 << xx[ii][2] << "   "
-  // 	 << vv[ii][0] << " "
-  // 	 << vv[ii][1] << " "
-  // 	 << vv[ii][2] << " "
-  // 	 << endl;
-  // }
-  // check variables
+  if (atom->map_style == Atom::MAP_NONE) {
+    error->all(FLERR, "Fix dplr requires an atom map, see atom_modify");
+  }
 
   if (xstr) {
     xvar = input->variable->find(xstr);
@@ -313,25 +304,61 @@ void FixDPLR::init() {
 
 /* ---------------------------------------------------------------------- */
 
+void FixDPLR::setup_post_neighbor() {
+  double **x = atom->x;
+
+  vector<pair<int, int> > valid_pairs;
+  get_valid_pairs(valid_pairs, true);
+
+  for (int ii = 0; ii < valid_pairs.size(); ++ii) {
+    int idx0 = valid_pairs[ii].first;
+    int idx1 = valid_pairs[ii].second;
+    int idx0_local = atom->map(atom->tag[idx0]);
+    int idx1_local = atom->map(atom->tag[idx1]);
+
+    for (int dd = 0; dd < 3; ++dd) {
+      x[idx1][dd] = x[idx0][dd];
+      x[idx0_local][dd] = x[idx0][dd];
+      x[idx1_local][dd] = x[idx0][dd];
+    }
+  }
+  int triclinic;
+  triclinic = domain->triclinic;
+  if (triclinic) {
+    domain->x2lamda(atom->nlocal);
+  }
+  domain->pbc();
+  domain->reset_box();
+  comm->setup();
+  neighbor->setup_bins();
+  comm->exchange();
+  comm->borders();
+  if (triclinic) {
+    domain->lamda2x(atom->nlocal + atom->nghost);
+  }
+
+  neighbor->build(1);
+}
+
+/* ---------------------------------------------------------------------- */
+
+void FixDPLR::setup_pre_exchange() {}
+
+/* ---------------------------------------------------------------------- */
+
 void FixDPLR::setup_pre_force(int vflag) { pre_force(vflag); }
 
 /* ---------------------------------------------------------------------- */
 
-void FixDPLR::setup(int vflag) {
-  // if (strstr(update->integrate_style,"verlet"))
-  post_force(vflag);
-  // else {
-  //   error->all(FLERR, "respa is not supported by this fix");
-  // }
-}
+void FixDPLR::setup(int vflag) { post_force(vflag); }
 
 /* ---------------------------------------------------------------------- */
 
 void FixDPLR::min_setup(int vflag) { setup(vflag); }
 
 /* ---------------------------------------------------------------------- */
 
-void FixDPLR::get_valid_pairs(vector<pair<int, int> > &pairs) {
+void FixDPLR::get_valid_pairs(vector<pair<int, int> > &pairs, bool is_setup) {
   pairs.clear();
 
   int nlocal = atom->nlocal;
@@ -398,7 +425,7 @@ void FixDPLR::get_valid_pairs(vector<pair<int, int> > &pairs) {
         error->all(FLERR, str);
       }
     }
-    if (!(idx0 < nlocal && idx1 < nlocal)) {
+    if (!(idx0 < nlocal && idx1 < nlocal) && (!is_setup)) {
       error->all(FLERR,
                  "find a bonded pair that is not on the same processor, "
                  "something should not happen");
@@ -409,7 +436,7 @@ void FixDPLR::get_valid_pairs(vector<pair<int, int> > &pairs) {
 
 /* ---------------------------------------------------------------------- */
 
-void FixDPLR::post_integrate() {
+void FixDPLR::pre_exchange() {
   double **x = atom->x;
   double **v = atom->v;
   int *type = atom->type;
@@ -418,7 +445,7 @@ void FixDPLR::post_integrate() {
   int nall = nlocal + nghost;
 
   vector<pair<int, int> > valid_pairs;
-  get_valid_pairs(valid_pairs);
+  get_valid_pairs(valid_pairs, false);
 
   for (int ii = 0; ii < valid_pairs.size(); ++ii) {
     int idx0 = valid_pairs[ii].first;
@@ -519,7 +546,7 @@ void FixDPLR::pre_force(int vflag) {
   // vector<int> & sort_fwd_map(atom_map.get_fwd_map());
 
   vector<pair<int, int> > valid_pairs;
-  get_valid_pairs(valid_pairs);
+  get_valid_pairs(valid_pairs, false);
 
   int odim = dpt.output_dim();
   assert(odim == 3);
@@ -642,7 +669,7 @@ void FixDPLR::post_force(int vflag) {
                                      list->firstneigh);
   // bonded pairs
   vector<pair<int, int> > valid_pairs;
-  get_valid_pairs(valid_pairs);
+  get_valid_pairs(valid_pairs, false);
   // output vects
   vector<FLOAT_PREC> dfcorr, dvcorr;
   // compute
@@ -728,7 +755,7 @@ void FixDPLR::post_force(int vflag) {
 
 /* ---------------------------------------------------------------------- */
 
-void FixDPLR::min_pre_exchange() { post_integrate(); }
+void FixDPLR::min_pre_exchange() { pre_exchange(); }
 
 /* ---------------------------------------------------------------------- */
 

source/lmp/fix_dplr.h

@@ -42,9 +42,11 @@ class FixDPLR : public Fix {
   int setmask() override;
   void init() override;
   void setup(int) override;
+  void setup_pre_exchange() override;
   void setup_pre_force(int) override;
+  void setup_post_neighbor() override;
   void min_setup(int) override;
-  void post_integrate() override;
+  void pre_exchange() override;
   void pre_force(int) override;
   void post_force(int) override;
   void min_pre_exchange() override;
@@ -72,7 +74,7 @@ class FixDPLR : public Fix {
   std::vector<double> efield;
   std::vector<double> efield_fsum, efield_fsum_all;
   int efield_force_flag;
-  void get_valid_pairs(std::vector<std::pair<int, int> > &pairs);
+  void get_valid_pairs(std::vector<std::pair<int, int> > &pairs, bool is_setup);
   int varflag;
   char *xstr, *ystr, *zstr;
   int xvar, yvar, zvar, xstyle, ystyle, zstyle;

source/lmp/tests/test_dplr.py

@@ -21,6 +21,7 @@
 dipole_pbtxt_file = Path(__file__).parent / "lrdipole.pbtxt"
 dipole_pb_file = Path(__file__).parent / "lrdipole.pb"
 data_file = Path(__file__).parent / "data.lmp"
+data_file2 = Path(__file__).parent / "data.lmp2"
 data_file_si = Path(__file__).parent / "data.si"
 data_type_map_file = Path(__file__).parent / "data_type_map.lmp"
 
@@ -241,6 +242,7 @@
 )
 
 box = np.array([0, 20, 0, 20, 0, 20, 0, 0, 0])
+box2 = np.array([0, 20, 0, 3.2575, 0, 20, 0, 0, 0])
 coord = np.array(
     [
         [1.25545000, 1.27562200, 0.98873000],
@@ -272,6 +274,9 @@ def setup_module():
     write_lmp_data_full(
         box, coord, mol_list, type_OH, charge, data_file, bond_list, mass_list
     )
+    write_lmp_data_full(
+        box2, coord, mol_list, type_OH, charge, data_file2, bond_list, mass_list
+    )
     write_lmp_data_full(
         box, coord, mol_list, type_HO, charge, data_type_map_file, bond_list, mass_list
     )
@@ -289,6 +294,7 @@ def setup_module():
 
 def teardown_module():
     os.remove(data_file)
+    os.remove(data_file2)
     os.remove(data_type_map_file)
     os.remove(data_file_si)
 
@@ -325,6 +331,13 @@ def lammps():
     lmp.close()
 
 
+@pytest.fixture
+def lammps2():
+    lmp = _lammps(data_file=data_file2)
+    yield lmp
+    lmp.close()
+
+
 @pytest.fixture
 def lammps_type_map():
     lmp = _lammps(data_file=data_type_map_file, exclude_type="2 3")
@@ -407,6 +420,20 @@ def test_pair_deepmd_lr(lammps):
     lammps.run(1)
 
 
+def test_pair_deepmd_lr_run0(lammps2):
+    lammps2.pair_style(f"deepmd {pb_file.resolve()}")
+    lammps2.pair_coeff("* *")
+    lammps2.bond_style("zero")
+    lammps2.bond_coeff("*")
+    lammps2.special_bonds("lj/coul 1 1 1 angle no")
+    lammps2.kspace_style("pppm/dplr 1e-5")
+    lammps2.kspace_modify(f"gewald {beta:.2f} diff ik mesh {mesh:d} {mesh:d} {mesh:d}")
+    lammps2.fix(f"0 all dplr model {pb_file.resolve()} type_associate 1 3 bond_type 1")
+    lammps2.fix_modify("0 virial yes")
+    lammps2.run(0)
+    lammps2.run(0)
+
+
 def test_pair_deepmd_lr_efield_constant(lammps):
     lammps.pair_style(f"deepmd {pb_file.resolve()}")
     lammps.pair_coeff("* *")