removal of residues i+4 and i-4 from calculation. only done in pyrose…

…tta version of DAN
hiranumn · Sep 14, 2022 · 27104c7 · 27104c7
1 parent cc3c191
commit 27104c7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 4 deletions.
diff --git a/deepAccNet/model.py b/deepAccNet/model.py
@@ -120,9 +120,13 @@ def calculate_LDDT(estogram, mask, center=7):
     # Get on the same device as indices
     device = estogram.device
 
-    # Remove diagonal from calculation
+    # Remove diagonal and residues i-4 to i+4 from calculation
     nres = mask.shape[-1]
-    mask = torch.mul(mask, torch.ones((nres, nres)).to(device) - torch.eye(nres).to(device))
+    diags = torch.ones((nres, nres)).to(device) - torch.eye(nres).to(device)
+    for i in range(1,5):
+        diags = diags - torch.diag(torch.ones(nres-i), diagonal=i)
+        diags = diags - torch.diag(torch.ones(nres-i), diagonal=-1*i)
+    mask = torch.mul(mask, diags)
     masked = torch.mul(estogram, mask)
 
     p0 = (masked[center]).sum(axis=0)

diff --git a/deepAccNet/model2.py b/deepAccNet/model2.py
@@ -289,9 +289,12 @@ def calculate_LDDT(estogram, mask, center=7):
     # Get on the same device as indices
     device = estogram.device
 
-    # Remove diagonal from calculation
+    # Remove diagonal and residues i-4 to i+4 from calculation
     nres = mask.shape[-1]
-    mask = torch.mul(mask, torch.ones((nres, nres)).to(device) - torch.eye(nres).to(device))
+    diags = torch.ones((nres, nres)).to(device) - torch.eye(nres).to(device)
+    for i in range(1,5):
+        diags = diags - torch.diag(torch.ones(nres-i), diagonal=i)
+        diags = diags - torch.diag(torch.ones(nres-i), diagonal=-1*i)
     masked = torch.mul(estogram, mask)
 
     p0 = (masked[center]).sum(axis=0)