From 4d257b00ee392fbffa37c7970ed8fc33466a46e5 Mon Sep 17 00:00:00 2001
From: sebastiankloth <kloth@mailbox.org>
Date: Thu, 3 Nov 2022 15:46:14 +0100
Subject: [PATCH] Added hbonds function

---
 mdevaluate/correlation.py  |  2 ++
 mdevaluate/distribution.py | 45 +++++++++++++++++++++++++++++++++++++-
 mdevaluate/utils.py        |  5 -----
 3 files changed, 46 insertions(+), 6 deletions(-)

diff --git a/mdevaluate/correlation.py b/mdevaluate/correlation.py
index 5f75072..ce50b3e 100644
--- a/mdevaluate/correlation.py
+++ b/mdevaluate/correlation.py
@@ -2,6 +2,8 @@ import numpy as np
 from scipy.special import legendre
 from itertools import chain
 import dask.array as darray
+from pathos.pools import ProcessPool
+from functools import partial
 
 from .autosave import autosave_data
 from .utils import filon_fourier_transformation, coherent_sum, histogram
diff --git a/mdevaluate/distribution.py b/mdevaluate/distribution.py
index ee4c048..a4bda18 100644
--- a/mdevaluate/distribution.py
+++ b/mdevaluate/distribution.py
@@ -4,7 +4,7 @@ from .coordinates import rotate_axis, polar_coordinates, spherical_coordinates
 from .atoms import next_neighbors
 from .autosave import autosave_data
 from .utils import runningmean
-from .pbc import pbc_diff, pbc_points
+from .pbc import pbc_diff, pbc_points, make_PBC
 from .logging import logger
 from scipy import spatial
 
@@ -357,3 +357,46 @@ def next_neighbor_distribution(atoms, reference=None, number_of_neighbors=4, bin
     resname_nn = reference.residue_names[nn]
     count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1)
     return np.histogram(count_nn, bins=bins, normed=normed)[0]
+    
+    
+def hbonds(D, H, A, box, DA_lim=0.35, HA_lim=0.35, min_cos=np.cos(30*np.pi/180), full_output=False):
+    def dist_DltA(D, H, A, box, max_dist=0.35):
+        ppoints, pind = make_PBC(D, box, thickness=max_dist+0.1, index=True)
+        Dtree = cKDTree(ppoints)
+        Atree = cKDTree(A)
+        pairs = Dtree.sparse_distance_matrix(Atree, max_dist, output_type='ndarray')
+        pairs = np.asarray(pairs.tolist())
+        pairs = np.int_(pairs[pairs[:,2] > 0][:,:2])
+        pairs[:,0] = pind[pairs[:,0]]
+        return pairs
+
+    def dist_AltD(D, H, A, box, max_dist=0.35):
+        ppoints, pind = make_PBC(A, box, thickness=max_dist+0.1, index=True)
+        Atree = cKDTree(ppoints)
+        Dtree = cKDTree(D)
+        pairs = Atree.sparse_distance_matrix(Dtree, max_dist, output_type='ndarray')   
+        pairs = np.asarray(pairs.tolist())
+        pairs = np.int_(pairs[pairs[:,2] > 0][:,:2])
+        pairs = pairs[:,::-1]
+        pairs[:,1] = pind[pairs[:,1]]
+        return pairs 
+  
+    if len(D) <= len(A):
+        pairs = dist_DltA(D,H,A,box,DA_lim)
+    else:
+        pairs = dist_AltD(D,H,A,box,DA_lim)
+    
+    vDH = md.pbc.pbc_diff(D[pairs[:,0]], H[pairs[:,0]], box)
+    vDA = md.pbc.pbc_diff(D[pairs[:,0]], A[pairs[:,1]], box)
+    vHA = md.pbc.pbc_diff(H[pairs[:,0]], A[pairs[:,1]], box)
+    angles_cos = np.clip(np.einsum('ij,ij->i', vDH, vDA)/
+          np.linalg.norm(vDH,axis=1)/
+          np.linalg.norm(vDA,axis=1), -1, 1)
+    is_bond = ((angles_cos >= min_cos) &
+                (np.sum(vHA**2,axis=-1) <= HA_lim**2) &
+                (np.sum(vDA**2,axis=-1) <= DA_lim**2))
+    if full_output:
+        return pairs[is_bond], angles_cos[is_bond], np.sum(vDA[is_bond]**2,axis=-1)**.5
+    else:
+        return pairs[is_bond]    
+
diff --git a/mdevaluate/utils.py b/mdevaluate/utils.py
index d539a1e..afaba2a 100644
--- a/mdevaluate/utils.py
+++ b/mdevaluate/utils.py
@@ -305,11 +305,6 @@ def Fqt_from_Grt(data, q):
     else:
         return isf.index, isf.values
 
-'''
-@numba.jit
-def norm(vec):
-    return (vec**2).sum()**0.5
-'''
 
 def singledispatchmethod(func):
     """A decorator to define a genric instance method, analogue to functools.singledispatch."""