From d375931ad664e16f3659af7c4b45d162b83f584f Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Sun, 7 Jul 2019 20:07:34 +0800
Subject: [PATCH] fix: apply PBC to bonds detected by OpenBabel (#62)

* Fixing style errors.

* fix: apply PBC to bonds detected by OpenBabel

fix #61

* fix self

* fix cell

* Update mddatasetbuilder/detect.py
---
 mddatasetbuilder/detect.py | 28 +++++++++++++++++++++++-----
 1 file changed, 23 insertions(+), 5 deletions(-)

diff --git a/mddatasetbuilder/detect.py b/mddatasetbuilder/detect.py
index 34c863e..979947b 100644
--- a/mddatasetbuilder/detect.py
+++ b/mddatasetbuilder/detect.py
@@ -7,6 +7,7 @@
 import openbabel
 import numpy as np
 from ase import Atom, Atoms
+from scipy.spatial import cKDTree
 from .dps import dps as connectmolecule
 
 
@@ -156,6 +157,16 @@ def readmolecule(self, lines):
     def _crd2bond(cls, step_atoms, readlevel):
         # copy from reacnetgenerator on 2019/4/13
         atomnumber = len(step_atoms)
+        if step_atoms.pbc.any():
+            # Apply period boundry conditions
+            # add ghost atoms
+            repeated_atoms = step_atoms.repeat(2)[atomnumber:]
+            tree = cKDTree(step_atoms.get_positions())
+            d = tree.query(repeated_atoms.get_positions(), k=1)[0]
+            nearest = d < 5
+            ghost_atoms = repeated_atoms[nearest]
+            realnumber = np.where(nearest)[0] % atomnumber
+            step_atoms += ghost_atoms
         xyzstring = ''.join((f"{atomnumber}\n{__name__}\n", "\n".join(
             [f'{s:2s} {x:22.15f} {y:22.15f} {z:22.15f}'
              for s, (x, y, z) in zip(
@@ -178,14 +189,21 @@ def _crd2bond(cls, step_atoms, readlevel):
                 if linecontent == 0:
                     s = line.split()
                     if len(s) > 3:
-                        b1, b2 = int(s[1])-1, int(s[2])-1
+                        s1, s2 = int(s[1])-1, int(s[2])-1
+                        if s1 >= atomnumber and s2 >= atomnumber:
+                            # duplicated
+                            continue
+                        elif s1 >= atomnumber:
+                            s1 = realnumber[s1-atomnumber]
+                        elif s2 >= atomnumber:
+                            s2 = realnumber[s2-atomnumber]
                         if readlevel:
                             level = 9 if s[3] == 'ar' else int(s[3])
-                            bondlevel[b1].append(level)
-                            bondlevel[b2].append(level)
+                            bondlevel[s1].append(level)
+                            bondlevel[s2].append(level)
                         else:
-                            bond[b1].append(b2)
-                            bond[b2].append(b1)
+                            bond[s1].append(s2)
+                            bond[s2].append(s1)
         return bondlevel if readlevel else bond
 
     def readcrd(self, item):