Some changes, atom -> point

src/Backboner.jl

@@ -1,7 +1,17 @@
 module Backboner
 
+using LinearAlgebra
+using Rotations: AngleAxis
+using NNlib: batched_mul
+
+norms(A::AbstractArray{<:Number}; dims=1) = sqrt.(sum(abs2, A; dims))
+dots(A1::AbstractArray{<:Number}, A2::AbstractMatrix{<:Number}; dims=1) = sum(A1 .* A2; dims)
+normalize_slices(A::AbstractArray{<:Number}; dims=1) = A ./ norms(A; dims)
+
+include("backbone.jl")
 export Backbone, coords
 
+include("bonds.jl")
 export ChainedBonds
 export append_bonds!, append_bonds
 export prepend_bonds!, prepend_bonds
@@ -10,14 +20,11 @@ export get_bond_lengths
 export get_bond_angles
 export get_torsion_angles
 
+include("frames.jl")
 export Frames
 
-export is_knotted
-
-include("backbone.jl")
-include("bonds.jl")
-include("frames.jl")
 include("knots.jl")
+export is_knotted
 
 using PrecompileTools
 

src/backbone.jl

@@ -1,7 +1,7 @@
 """
     Backbone{T<:Real,M<:AbstractMatrix{T}} <: AbstractVector{AbstractVector{T}}
 
-The `Backbone` type is designed to efficiently store and manipulate the three-dimensional coordinates of backbone atoms.
+The `Backbone` type is designed to efficiently store and manipulate the three-dimensional coordinates of backbone points.
 """
 struct Backbone{T<:Real,M<:AbstractMatrix{T}} <: AbstractVector{AbstractVector{T}}
     coords::M

src/bonds.jl

@@ -1,17 +1,24 @@
-using LinearAlgebra
-using Rotations: AngleAxis
-
-norms(A::AbstractArray{<:Real}; dims=1) = sqrt.(sum(abs2, A; dims))
-dots(A1::AbstractArray{T}, A2::AbstractMatrix{T}; dims=1) where T <: Real = sum(A1 .* A2; dims)
-normalize_slices(A::AbstractArray{<:Real}; dims=1) = A ./ norms(A; dims)
-
 # TODO: see if views would be differentiable (GPU?)
 
-get_atom_displacements(backbone::Backbone, start::Integer, step::Integer, stride::Integer) =
+"""
+    get_displacements(backbone, start, step, stride)
+
+Get the displacements between points in `backbone`, where `start` is the index of the first point,
+`step` is the offset of the other points along the backbone (e.g. `1` for bond vectors),
+and `stride` is the number of points to skip after each step.
+"""
+get_displacements(backbone::Backbone, start::Integer, step::Integer, stride::Integer) =
     backbone.coords[:, start+step:stride:end] .- backbone.coords[:, start:stride:end-step]
 
-get_atom_distances(backbone::Backbone, start::Integer, step::Integer, stride::Integer) =
-    norms(get_atom_displacements(backbone, start, step, stride))
+"""
+    get_distances(backbone, start, step, stride)
+
+Get the Euclidean distances between points in `backbone`, where `start` is the index of the first point,
+`step` is the offset of the other points along the backbone (e.g. `1` for bond lengths),
+and `stride` is the number of points to skip after each step.
+"""
+get_distances(backbone::Backbone, start::Integer, step::Integer, stride::Integer) =
+    norms(get_displacements(backbone, start, step, stride))
 
 _get_bond_lengths(bond_vectors::AbstractMatrix{<:Real}) = norms(bond_vectors)
 
@@ -39,32 +46,20 @@ function _get_torsion_angles(bond_vectors::AbstractMatrix{<:Real})
     return torsion_angles
 end
 
-get_bond_vectors(backbone::Backbone) = get_atom_displacements(backbone, 1, 1, 1)
+get_bond_vectors(backbone::Backbone) = get_displacements(backbone, 1, 1, 1)
 get_bond_lengths(backbone::Backbone) = _get_bond_lengths(get_bond_vectors(backbone)) |> vec
 get_bond_angles(backbone::Backbone) = _get_bond_angles(get_bond_vectors(backbone)) |> vec
 get_torsion_angles(backbone::Backbone) = _get_torsion_angles(get_bond_vectors(backbone)) |> vec
 
 """
-    ChainedBonds{T <: Real, V <: AbstractVector{T}}
+    ChainedBonds{T<:Real,V<:AbstractVector{T}}
 
 A lazy way to store a backbone as a series of bond lengths, bond angles, and torsion_angles.
 
 # Examples
 
 ```jldoctest
-julia> backbone = Protein.readpdb("test/data/1ZAK.pdb")["A"].backbone
-660-element Backbone{Float64, Matrix{Float64}}:
- [22.346, 17.547, 23.294]
- [22.901, 18.031, 21.993]
- [23.227, 16.793, 21.163]
- [24.115, 16.923, 20.175]
- [24.478, 15.779, 19.336]
- ⋮
- [21.480, 14.668, 4.974]
- [22.041, 14.866, 3.569]
- [21.808, 13.861, 2.734]
- [22.263, 13.862, 1.355]
- [21.085, 14.233, 0.446]
+julia> backbone = Backbone(randn(Float32, 3, 660));
 
 julia> bonds = ChainedBonds(backbone)
 ChainedBonds{Float64, Vector{Float64}} with 659 bond lengths, 658 bond angles, and 657 torsion angles
@@ -129,7 +124,7 @@ function get_backbone_start(bonds::ChainedBonds{T}) where T
     return Backbone(coords)
 end
 
-function get_next_atom(bonds, i, A, B, C)
+function get_next_point(bonds, i, A, B, C)
     AB, BC = B - A, C - B
     n_AB, n_BC = normalize(AB), normalize(BC)
     CD_init = n_BC * bonds.bond_lengths[i-1]
@@ -156,7 +151,7 @@ function Backbone(
     @assert 3 <= l <= L
     coords[:, 1:l] = backbone_start.coords
     for i in l+1:L
-        coords[:, i] = get_next_atom(bonds, i, eachcol(coords[:, i-3:i-1])...)
+        coords[:, i] = get_next_point(bonds, i, eachcol(coords[:, i-3:i-1])...)
     end
     return Backbone(coords)
 end
@@ -177,12 +172,12 @@ end
     append_bonds(backbone, bond_lengths, bond_angles, torsion_angles)
 """
 function append_bonds(backbone::Backbone, bond_lengths::AbstractVector{<:Real}, bond_angles::AbstractVector{<:Real}, torsion_angles::AbstractVector{<:Real})
-    length(backbone) >= 3 || throw(ArgumentError("backbone must have at least 3 atoms"))
+    length(backbone) >= 3 || throw(ArgumentError("backbone must have at least 3 points"))
     length(bond_lengths) == length(bond_angles) == length(torsion_angles) || throw(ArgumentError("bond_lengths, bond_angles, and torsion_angles must have the same length"))
-    last_three_atoms_backbone = backbone[end-2:end]
-    bonds_end = ChainedBonds(last_three_atoms_backbone)
+    last_three_points_backbone = backbone[end-2:end]
+    bonds_end = ChainedBonds(last_three_points_backbone)
     append_bonds!(bonds_end, bond_lengths, bond_angles, torsion_angles)
-    backbone_end = Backbone(bonds_end, backbone_start=last_three_atoms_backbone)
+    backbone_end = Backbone(bonds_end, backbone_start=last_three_points_backbone)
     new_backbone = Backbone(cat(backbone.coords, backbone_end[4:end].coords, dims=2))
     return new_backbone
 end
@@ -202,12 +197,12 @@ end
     prepend_bonds(backbone, bond_lengths, bond_angles, torsion_angles)
 """
 function prepend_bonds(backbone::Backbone, bond_lengths::AbstractVector{<:Real}, bond_angles::AbstractVector{<:Real}, torsion_angles::AbstractVector{<:Real})
-    length(backbone) >= 3 || throw(ArgumentError("backbone must have at least 3 atoms"))
-    first_three_atoms_backbone = backbone[1:3]
-    bonds_start = ChainedBonds(first_three_atoms_backbone)
+    length(backbone) >= 3 || throw(ArgumentError("backbone must have at least 3 points"))
+    first_three_points_backbone = backbone[1:3]
+    bonds_start = ChainedBonds(first_three_points_backbone)
     prepend_bonds!(bonds_start, bond_lengths, bond_angles, torsion_angles)
     reverse!(bonds_start)
-    backbone_start = Backbone(bonds_start, backbone_start=first_three_atoms_backbone[3:-1:1])
+    backbone_start = Backbone(bonds_start, backbone_start=first_three_points_backbone[3:-1:1])
     new_backbone = Backbone(cat(backbone_start[end:-1:4].coords, backbone.coords, dims=2))
     return new_backbone
 end

src/frames.jl

@@ -1,6 +1,3 @@
-using LinearAlgebra
-using NNlib
-
 centroid(A::AbstractArray{<:Real}; dims=2) = sum(A; dims) ./ prod(size(A)[dims])
 
 #= FIXME: P currently gets aligned to Q, should be the other way around?
@@ -53,7 +50,7 @@ function Frames(backbone::Backbone{T}, ideal_coords::AbstractMatrix{<:Real}) whe
     return Frames(rotations, translations)
 end
 
-(frames::Frames{T})(coords::AbstractMatrix{T}) where T<:Real = frames.rotations ⊠ (coords .- centroid(coords)) .+ reshape(frames.translations, 3, 1, :)
+(frames::Frames{T})(coords::AbstractMatrix{T}) where T<:Real = batched_mul(frames.rotations, (coords .- centroid(coords))) .+ reshape(frames.translations, 3, 1, :)
 (frames::Frames{T})(coords::AbstractMatrix{<:Real}) where T<:Real = frames(T.(coords))
 
 Backbone(frames::Frames, ideal_coords::AbstractMatrix{<:Real}) = Backbone(frames(ideal_coords))
@@ -116,18 +113,19 @@ function rotation_matrices_to_quaternions(R::AbstractArray{<:Real,3})
           1 .- r11 - r22 + r33]
 
     # 4x4xN
-    Q = hcat(q0, q1, q2, q3)
+    qs = hcat(q0, q1, q2, q3)
 
     # 1x4xN, norm of each quaternion
-    norms = sqrt.(sum(abs2, Q, dims=1))
+    q_norms = norms(qs, dims=1)
+
+    exp_norms = exp.(q_norms)
 
-    exp_norms = exp.(norms)
     # 1x4xN, norm weights
     weights = exp_norms ./ sum(exp_norms, dims=3)
 
     # batched matmul, 4x1xN
-    q = Q ⊠ reshape(weights, 4, 1, :)
-    q_normalized = q ./ sqrt.(sum(abs2, q, dims=1))
+    Q = batched_mul(qs, reshape(weights, 4, 1, :))
+    Q_normalized = q ./ norms(Q, dims=1)
 
     return reshape(q_normalized, 4, :)
-end
+end

test/bonds.jl

@@ -8,14 +8,14 @@
             0.0 2.0 4.0 4.0
         ])
 
-        @testset "get_atom_displacements" begin
-            @test Backboner.get_atom_displacements(backbone, 1, 1, 2) == [1.0 3.0; 2.0 4.0; 2.0 0.0]
-            @test Backboner.get_atom_displacements(backbone, 2, 1, 2) == [1.0; 2.0; 2.0;;]
+        @testset "get_displacements" begin
+            @test Backboner.get_displacements(backbone, 1, 1, 2) == [1.0 3.0; 2.0 4.0; 2.0 0.0]
+            @test Backboner.get_displacements(backbone, 2, 1, 2) == [1.0; 2.0; 2.0;;]
         end
 
-        @testset "get_atom_distances" begin
-            @test Backboner.get_atom_distances(backbone, 1, 1, 2) == [3.0 5.0]
-            @test Backboner.get_atom_distances(backbone, 2, 1, 2) == [3.0;;]
+        @testset "get_distances" begin
+            @test Backboner.get_distances(backbone, 1, 1, 2) == [3.0 5.0]
+            @test Backboner.get_distances(backbone, 2, 1, 2) == [3.0;;]
         end
 
         @testset "get_bond_vectors" begin