From 486f7228a79dfb3fecbdd913226f862a24677671 Mon Sep 17 00:00:00 2001 From: anton083 Date: Fri, 13 Sep 2024 09:37:11 +0200 Subject: [PATCH] Some changes, atom -> point --- src/Backboner.jl | 17 +++++++++---- src/backbone.jl | 2 +- src/bonds.jl | 65 ++++++++++++++++++++++-------------------------- src/frames.jl | 18 ++++++-------- test/bonds.jl | 12 ++++----- 5 files changed, 57 insertions(+), 57 deletions(-) diff --git a/src/Backboner.jl b/src/Backboner.jl index 0692064..13e37bf 100644 --- a/src/Backboner.jl +++ b/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 diff --git a/src/backbone.jl b/src/backbone.jl index 61d1188..e21b9c7 100644 --- a/src/backbone.jl +++ b/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 diff --git a/src/bonds.jl b/src/bonds.jl index 585c564..eabd783 100644 --- a/src/bonds.jl +++ b/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 \ No newline at end of file diff --git a/src/frames.jl b/src/frames.jl index 00f4c7f..e4d6e4c 100644 --- a/src/frames.jl +++ b/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 \ No newline at end of file +end diff --git a/test/bonds.jl b/test/bonds.jl index 138fc45..0e1d3f3 100644 --- a/test/bonds.jl +++ b/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