Use KDTree

MurrellGroup · Nov 7, 2024 · a545ba0 · a545ba0
1 parent 845d33c
commit a545ba0
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,17 +1,19 @@
 name = "AssigningSecondaryStructure"
 uuid = "8ed43e74-60fb-4e11-99b9-91deed37aef7"
 authors = ["Anton Oresten <anton.oresten42@gmail.com>"]
-version = "0.5.0"
+version = "0.5.1"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NearestNeighbors = "b8a86587-4115-5ab1-83bc-aa920d37bbce"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
 LinearAlgebra = "1"
 NearestNeighbors = "0.4.20"
 SparseArrays = "1.11.0"
+StaticArrays = "1.9.8"
 julia = "1"
 
 [extras]

diff --git a/src/hydrogen_bonds.jl b/src/hydrogen_bonds.jl
@@ -31,11 +31,12 @@ end
 
 using NearestNeighbors
 using SparseArrays
+using StaticArrays
 
 function get_hbond_map(coords::Array{T,3}, cutoff::Real=CUTOFF) where T<:Real
-    C_pos = coords[:, 3, :]
+    C_pos = @views coords[:, 3, :]
     O_pos = get_oxygen_positions(coords)
-    N_pos = coords[:, 1, :]
+    N_pos = @views coords[:, 1, :]
     H_pos = get_hydrogen_positions(coords)
 
     num_residues = size(coords, 3)
@@ -44,18 +45,23 @@ function get_hbond_map(coords::Array{T,3}, cutoff::Real=CUTOFF) where T<:Real
     # Build KD-tree for N positions
     N_tree = KDTree(N_pos)
 
+    C_pos_static = SVector{3, T}.(eachslice(C_pos, dims=2))
+    O_pos_static = SVector{3, T}.(eachslice(O_pos, dims=2))
+    N_pos_static = SVector{3, T}.(eachslice(N_pos, dims=2))
+    H_pos_static = SVector{3, T}.(eachslice(H_pos, dims=2))
+
     # For each O_i, find N_j within cutoff
     for i in 1:num_residues
-        O_i = O_pos[:, i]
-        idxs = inrange(N_tree, O_i, 10.0)
+        O_i = O_pos_static[i]
+        idxs = inrange(N_tree, O_i, 5.0)
         for j in idxs
             # Exclude self and neighboring residues
             if abs(i - j) > 2  # Exclude i == j, i+1, i+2
                 # Compute E[i,j] as before
-                ON_dist = norm(O_pos[:, i] - N_pos[:, j])
-                CH_dist = norm(C_pos[:, i] - H_pos[:, j])
-                OH_dist = norm(O_pos[:, i] - H_pos[:, j])
-                CN_dist = norm(C_pos[:, i] - N_pos[:, j])
+                ON_dist = norm(O_pos_static[i] - N_pos_static[j])
+                CH_dist = norm(C_pos_static[i] - H_pos_static[j])
+                OH_dist = norm(O_pos_static[i] - H_pos_static[j])
+                CN_dist = norm(C_pos_static[i] - N_pos_static[j])
                 E_ij = T(Q1Q2 * F) * (1 / ON_dist + 1 / CH_dist - 1 / OH_dist - 1 / CN_dist)
                 if E_ij < cutoff
                     Hbonds[i, j] = true