Coordination number for RDF and PBC for water molecules.

dpdata/md/rdf.py

@@ -21,6 +21,15 @@ def rdf(sys,
         Maximal range of rdf calculation
     nbins: int
         Number of bins for rdf calculation
+
+    Returns
+    -------
+    xx: np.array
+        The lattice of r
+    rdf: np.array
+        The value of rdf at r
+    coord: np.array
+        The coordination number up to r
     """
     return compute_rdf(sys['cells'], sys['coords'], sys['atom_types'], 
                        sel_type = sel_type, 
@@ -36,12 +45,16 @@ def compute_rdf(box,
     nframes = box.shape[0]
     xx = None
     all_rdf = []
+    all_cod = []
     for ii in range(nframes):
-        xx, rdf = _compute_rdf_1frame(box[ii], posis[ii], atype, sel_type, max_r, nbins)
+        xx, rdf, cod = _compute_rdf_1frame(box[ii], posis[ii], atype, sel_type, max_r, nbins)
         all_rdf.append(rdf)
+        all_cod.append(cod)
     all_rdf = np.array(all_rdf).reshape([nframes, -1])
+    all_cod = np.array(all_cod).reshape([nframes, -1])
     all_rdf = np.average(all_rdf, axis = 0)
-    return xx, all_rdf
+    all_cod = np.average(all_cod, axis = 0)
+    return xx, all_rdf, all_cod
 
 def _compute_rdf_1frame(box, 
                         posis, 
@@ -65,6 +78,7 @@ def _compute_rdf_1frame(box,
     nlist = ase.neighborlist.NeighborList(max_r, self_interaction=False, bothways=True, primitive=ase.neighborlist.NewPrimitiveNeighborList)
     nlist.update(atoms)
     stat = np.zeros(nbins)
+    stat_acc = np.zeros(nbins)
     hh = max_r / float(nbins)
     for ii in range(natoms) :
         # atom "0"
@@ -90,13 +104,17 @@ def _compute_rdf_1frame(box,
     for ii in sel_type[1]:
         c1 += np.sum(atype == ii)
     rho1 = c1 / np.linalg.det(box)
+    # compute coordination number
+    for ii in range(1, nbins):
+        stat_acc[ii] = stat_acc[ii-1] + stat[ii-1]
+    stat_acc = stat_acc / c0
     # compute rdf
     for ii in range(nbins):
         vol = 4./3. * np.pi * ( ((ii+1)*hh) ** 3 - ((ii)*hh) ** 3 )
         rho = stat[ii] / vol
         stat[ii] = rho / rho1 / c0
     xx = np.arange(0, max_r-1e-12, hh)
-    return xx, stat
+    return xx, stat, stat_acc
 
 if __name__ == '__main__':
     import dpdata

dpdata/md/water.py

@@ -1,5 +1,6 @@
 import numpy as np
 from .pbc import posi_diff
+from .pbc import posi_shift
 
 def compute_bonds (box, 
                    posis,
@@ -179,3 +180,26 @@ def find_ions (atype,
             else :
                 raise RuntimeError("unknow case: numb of O bonded to H > 1")
     return no, noh, noh2, noh3, nh
+
+
+
+def pbc_coords(cells, 
+               coords, 
+               atom_types, 
+               oh_sel = [0, 1],
+               max_roh = 1.3):
+    bonds = compute_bonds(cells, coords, atom_types, oh_sel = oh_sel, max_roh = max_roh, uniq_hbond = True)
+
+    new_coords = np.copy(coords)
+    natoms = len(atom_types)
+    o_type = oh_sel[0]
+    h_type = oh_sel[1]
+    for ii in range(natoms):
+        if atom_types[ii] == o_type:
+            assert(len(bonds[ii]) == 2), 'O has more than 2 bonded Hs, stop'
+            for jj in bonds[ii]:
+                assert(atom_types[jj] == h_type), 'The atom bonded to O is not H, stop'
+                shift = posi_shift(cells, coords[jj], coords[ii])
+                new_coords[jj] = coords[jj] + np.matmul(shift, cells)
+    return new_coords
+