Merge pull request #25 from m3g/coordination_number

add coordination_number function

docs/src/miscelaneous.md

@@ -8,6 +8,7 @@ CollapsedDocStrings = true
 Modules = [ MolSimToolkit ]
 Pages = [ 
     "bulk_coordination.jl",
+    "coordination_number.jl", 
     "intermittent_correlation.jl",
     "distances.jl", 
     "center_of_mass.jl",

docs/src/molecular_minimum_distances.md

@@ -81,6 +81,11 @@ julia> count(x -> x.within_cutoff, list)
 
 Thus, 33 TMAO molecules are within the cutoff distance from the protein, and the distances can be used to study the solvation of the protein.
 
+!!! tip
+    The `coordination_number` function of this package essentially performs the above calculation iteratively 
+    along a trajectory. The source of of such function is simple and can be used to further understand the utility
+    and usage of the minimum-distance calculations.
+
 ## Performance
 
 This package exists because this computation is fast. For example, let us choose the water molecules instead, and benchmark the time required to compute this set of distances:

src/MolSimToolkit.jl

@@ -9,7 +9,7 @@ using TestItems: @testitem
 using StaticArrays: FieldVector, SMatrix, MVector
 using LinearAlgebra: norm, cross, dot
 using Reexport: @reexport
-using ProgressMeter: Progress, next!
+using ProgressMeter: Progress, next!, @showprogress
 using Statistics: mean
 
 export wrap, wrap_to_first
@@ -18,6 +18,7 @@ export dihedral, dihedrals, average_dihedrals
 export align, align!, rmsd, rmsd_matrix
 export intermittent_correlation
 export bulk_coordination
+export coordination_number
 
 # Reexported from ProteinSecondaryStructures for convenience
 using ProteinSecondaryStructures: dssp_run, stride_run, 
@@ -53,6 +54,7 @@ include("./miscelaneous_functions/center_of_mass.jl")
 include("./miscelaneous_functions/intermittent_correlation.jl")
 include("./miscelaneous_functions/bulk_coordination.jl")
 include("./miscelaneous_functions/most_representative_structure.jl")
+include("./miscelaneous_functions/coordination_number.jl")
 
 #  Structural alignment
 include("./procrustes.jl")

src/miscelaneous_functions/coordination_number.jl

@@ -0,0 +1,102 @@
+"""
+    coordination_number(
+        sim::Simulation,
+        solute::AbstractVector{<:PDBTools.Atom},
+        solvent::AbstractVector{<:PDBTools.Atom};
+        solvent_natomspermol::Int,
+        cutoff::Real,
+        show_progress::Bool = true
+    )
+
+Calculate the coordination number of the solute atoms with the solvent atoms.
+
+# Positional Arguments
+
+- `sim::Simulation`: Simulation object.
+- `solute::AbstractVector{<:PDBTools.Atom}`: Vector of solute atoms.
+- `solvent::AbstractVector{<:PDBTools.Atom}`: Vector of solvent atoms.
+
+# Keyword Arguments
+
+- `solvent_natomspermol::Int`: Number of atoms per solvent molecule.
+- `cutoff::Real`: Cutoff distance.
+- `show_progress::Bool`: Show progress bar. (optional, default: `true`)
+
+# Returns
+
+- `cn::Vector{Float64}`: Vector with the coordination number of the solute atoms with the solvent atoms, at each frame.
+
+# Example
+
+```jldoctest
+julia> using MolSimToolkit, PDBTools, MolSimToolkit.Testing
+
+julia> sim = Simulation(Testing.namd2_pdb, Testing.namd2_traj; frames=1:5);
+
+julia> protein = select(atoms(sim), "protein");
+
+julia> tmao = select(atoms(sim), "resname TMAO");
+
+julia> coordination_number(sim, protein, tmao; solvent_natomspermol=14, cutoff=3.0, show_progress=false)
+5-element Vector{Int64}:
+ 7
+ 3
+ 4
+ 5
+ 6
+
+```
+
+"""
+function coordination_number(
+    sim::Simulation,
+    solute::AbstractVector{<:PDBTools.Atom},
+    solvent::AbstractVector{<:PDBTools.Atom};
+    solvent_natomspermol::Int,
+    cutoff::Real,
+    show_progress::Bool = true,
+)
+    if length(solvent) % solvent_natomspermol != 0
+        throw(ArgumentError("""\n
+            Number of solvent atoms is not a multiple of solvent_natomspermol
+
+        """))
+    end
+    inds_solute = PDBTools.index.(solute)
+    inds_solvent = PDBTools.index.(solvent)
+    cn = zeros(Int, length(sim))
+    first_frame!(sim)
+    p = positions(current_frame(sim))
+    sys = CrossPairs(
+        xpositions=p[inds_solvent],
+        ypositions=p[inds_solute],
+        xn_atoms_per_molecule=solvent_natomspermol,
+        cutoff=cutoff,
+        unitcell=unitcell(current_frame(sim)),
+    )
+    prg = Progress(length(sim); enabled=show_progress)
+    for (iframe, frame) in enumerate(sim)
+        p = positions(frame)
+        sys.xpositions .= @view(p[inds_solvent])
+        sys.ypositions .= @view(p[inds_solute])
+        sys.unitcell = unitcell(frame)
+        md_list = minimum_distances!(sys)
+        cn[iframe] = count(md -> md.within_cutoff, md_list)
+        next!(prg)
+    end
+    return cn
+end
+
+@testitem "coordination_number" begin
+    using MolSimToolkit, PDBTools, MolSimToolkit.Testing
+    sim = Simulation(Testing.namd2_pdb, Testing.namd2_traj)
+    protein = select(atoms(sim), "protein")
+    tmao = select(atoms(sim), "resname TMAO")
+    cn = coordination_number(
+        sim, protein, tmao; 
+        solvent_natomspermol=14, cutoff=3.0,
+        show_progress=false
+    )
+    @test cn == [7, 3, 4, 5, 6, 5, 7, 7, 10, 5, 4, 4, 5, 6, 2, 9, 4, 5, 2, 6]
+end
+