From 62aebb93023c4f471b7b9741a7f28c3992c0c518 Mon Sep 17 00:00:00 2001
From: Han Wang <wang_han@iapcm.ac.cn>
Date: Fri, 3 Jan 2020 17:14:13 +0800
Subject: [PATCH] use nlist built by ase

---
 dpdata/md/water.py       | 76 ++++++++++++++++++++++++++++++++++++----
 tests/test_water_ions.py | 27 +++++++++++++-
 2 files changed, 96 insertions(+), 7 deletions(-)

diff --git a/dpdata/md/water.py b/dpdata/md/water.py
index afc171a3d..8c38a108a 100644
--- a/dpdata/md/water.py
+++ b/dpdata/md/water.py
@@ -1,12 +1,76 @@
 import numpy as np
 from .pbc import posi_diff
 
-def compute_bonds(box,
-                  posis, 
-                  atype, 
-                  oh_sel = [0,1],
-                  max_roh = 1.3, 
-                  uniq_hbond = True):
+def compute_bonds (box, 
+                   posis,
+                   atype, 
+                   oh_sel = [0,1],
+                   max_roh = 1.3, 
+                   uniq_hbond = True):
+    try :
+        import ase
+        import ase.neighborlist    
+        # nlist implemented by ase
+        return compute_bonds_ase(box, posis, atype, oh_sel, max_roh, uniq_hbond)        
+    except ImportError:
+        # nlist naivly implemented , scales as O(N^2)
+        return compute_bonds_naive(box, posis, atype, oh_sel, max_roh, uniq_hbond)
+
+
+def compute_bonds_ase(box, 
+                      posis,
+                      atype, 
+                      oh_sel = [0,1],
+                      max_roh = 1.3, 
+                      uniq_hbond = True):
+    natoms = len(posis)
+    from ase import Atoms
+    import ase.neighborlist    
+    atoms = Atoms(positions=posis, cell=box, pbc=[1,1,1])
+    nlist = ase.neighborlist.NeighborList(max_roh, self_interaction=False, bothways=True, primitive=ase.neighborlist.NewPrimitiveNeighborList)
+    nlist.update(atoms)
+    bonds = []
+    o_type = oh_sel[0]
+    h_type = oh_sel[1]
+    for ii in range(natoms) :
+        bonds.append([])
+    for ii in range(natoms) :
+        if atype[ii] == o_type :
+            nn, ss = nlist.get_neighbors(ii)
+            for jj in nn:
+                if atype[jj] == h_type :
+                    dr = posi_diff(box, posis[ii], posis[jj])
+                    if np.linalg.norm(dr) < max_roh :
+                        bonds[ii].append(jj)
+                        bonds[jj].append(ii)
+    if uniq_hbond :
+        for jj in range(natoms) :
+            if atype[jj] == h_type :
+                if len(bonds[jj]) > 1 :
+                    orig_bonds = bonds[jj]
+                    min_bd = 1e10
+                    min_idx = -1
+                    for ii in bonds[jj] :
+                        dr = posi_diff(box, posis[ii], posis[jj])
+                        drr = np.linalg.norm(dr)
+                        # print(ii,jj, posis[ii], posis[jj], drr)
+                        if drr < min_bd :
+                            min_idx = ii
+                            min_bd = drr
+                    bonds[jj] = [min_idx]
+                    orig_bonds.remove(min_idx)
+                    # print(min_idx, orig_bonds, bonds[jj])
+                    for ii in orig_bonds :
+                        bonds[ii].remove(jj)
+    return bonds
+                      
+
+def compute_bonds_naive(box,
+                        posis, 
+                        atype, 
+                        oh_sel = [0,1],
+                        max_roh = 1.3, 
+                        uniq_hbond = True):
     natoms = len(posis)
     bonds = []
     o_type = oh_sel[0]
diff --git a/tests/test_water_ions.py b/tests/test_water_ions.py
index 7027331c4..3a9436b6d 100644
--- a/tests/test_water_ions.py
+++ b/tests/test_water_ions.py
@@ -2,6 +2,12 @@
 import numpy as np
 import unittest
 from context import dpdata
+try:
+    import ase
+    import ase.neighborlist    
+    exist_ase=True
+except:
+    exist_ase=False
 
 class TestIons(unittest.TestCase):
     
@@ -13,7 +19,6 @@ def setUp(self):
                                                    self.system.data['coords'][0],
                                                    self.system.data['atom_types'])
 
-
     def test_ions_count(self) :
         no, noh, noh2, noh3, nh \
             = dpdata.md.water.find_ions(self.system.data['atom_types'], self.bonds)
@@ -25,6 +30,26 @@ def test_ions_count(self) :
         self.assertEqual(noh[0], 0)
         self.assertEqual(noh3[0], 51)
 
+
+@unittest.skipIf(not exist_ase, "skip TestAseComputeBond")
+class TestAseComputeBond(unittest.TestCase):
+    def setUp(self):
+        self.system = dpdata.System()
+        self.system.from_lammps_lmp(os.path.join('poscars', 'conf.waterion.lmp'), 
+                                    type_map = ['O', 'H'])
+        self.bonds = dpdata.md.water.compute_bonds_naive(self.system.data['cells'][0],
+                                                         self.system.data['coords'][0],
+                                                         self.system.data['atom_types'])
+        self.bonds_ase = dpdata.md.water.compute_bonds_ase(self.system.data['cells'][0],
+                                                           self.system.data['coords'][0],
+                                                           self.system.data['atom_types'])
+
+    def test_bond_identity(self):
+        self.assertTrue(len(self.bonds), len(self.bonds_ase))
+        for ii in range(len(self.bonds)):
+            self.assertTrue(set(self.bonds[ii]) == set(self.bonds_ase[ii]))
+
+
         
 if __name__ == '__main__':
     unittest.main()