support rdf

dpdata/md/rdf.py

@@ -0,0 +1,106 @@
+import numpy as np
+
+def rdf(sys,
+        sel_type = [None, None], 
+        max_r = 5, 
+        nbins = 100) :
+    """
+    compute the rdf of a system
+
+    Parameters
+    ----------
+    sys : System or LabeledSystem
+        The dpdata system
+    sel_type: list
+        List of size 2. The first element specifies the type of the first atom, 
+        while the second element specifies the type of the second atom. 
+        Both elements can be ints or list of ints. 
+        If the element is None, all types are specified.
+        Examples are sel_type = [0, 0], sel_type = [0, [0, 1]] or sel_type = [0, None]
+    max_r: float
+        Maximal range of rdf calculation
+    nbins: int
+        Number of bins for rdf calculation
+    """
+    return compute_rdf(sys['cells'], sys['coords'], sys['atom_types'], 
+                       sel_type = sel_type, 
+                       max_r = max_r, 
+                       nbins = nbins)
+
+def compute_rdf(box,
+                posis, 
+                atype, 
+                sel_type = [None, None], 
+                max_r = 5, 
+                nbins = 100) :
+    nframes = box.shape[0]
+    xx = None
+    all_rdf = []
+    for ii in range(nframes):
+        xx, rdf = _compute_rdf_1frame(box[ii], posis[ii], atype, sel_type, max_r, nbins)
+        all_rdf.append(rdf)
+    all_rdf = np.array(all_rdf).reshape([nframes, -1])
+    all_rdf = np.average(all_rdf, axis = 0)
+    return xx, all_rdf
+
+def _compute_rdf_1frame(box, 
+                        posis, 
+                        atype, 
+                        sel_type = [None, None], 
+                        max_r = 5, 
+                        nbins = 100) :
+    all_types = list(set(list(np.sort(atype))))
+    if sel_type[0] is None:
+        sel_type[0] = all_types
+    if sel_type[1] is None:
+        sel_type[1] = all_types
+    if type(sel_type[0]) is not list:
+        sel_type[0] = [sel_type[0]]
+    if type(sel_type[1]) is not list:
+        sel_type[1] = [sel_type[1]]
+    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_r, self_interaction=False, bothways=True, primitive=ase.neighborlist.NewPrimitiveNeighborList)
+    nlist.update(atoms)
+    stat = np.zeros(nbins)
+    hh = max_r / float(nbins)
+    for ii in range(natoms) :
+        # atom "0"
+        if atype[ii] in sel_type[0]:
+            indices, offsets = nlist.get_neighbors(ii)
+            for jj, os in zip(indices, offsets):
+                # atom "1"
+                if atype[jj] in sel_type[1]:
+                    posi_jj = atoms.positions[jj] + np.dot(os, atoms.get_cell())
+                    diff = posi_jj - atoms.positions[ii]
+                    dr = np.linalg.norm(diff)                    
+                    # if (np.linalg.norm(diff- diff_1)) > 1e-12 :
+                    #     raise RuntimeError
+                    si = int(dr / hh)
+                    if si < nbins:
+                        stat[si] += 1
+    # count the number of atom1
+    c0 = 0
+    for ii in sel_type[0]:
+        c0 += np.sum(atype == ii)
+    # count the number of atom1
+    c1 = 0
+    for ii in sel_type[1]:
+        c1 += np.sum(atype == ii)
+    rho1 = c1 / np.linalg.det(box)
+    # 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
+
+if __name__ == '__main__':
+    import dpdata
+    sys = dpdata.System('out.lmp')
+    xx, stat = rdf(sys, sel_type = [[0], None], max_r = 8, nbins = 100)
+    res = np.concatenate([xx, stat]).reshape([2, -1])
+    np.savetxt('rdf.out', res.T)