some changes

Project.toml

@@ -4,12 +4,13 @@ authors = ["Anton Oresten"]
 version = "1.0.0-DEV"
 
 [deps]
+ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
+Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
 Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
-GR = "28b8d3ca-fb5f-59d9-8090-bfdbd6d07a71"
+GLMakie = "e9467ef8-e4e7-5192-8a1a-b1aee30e663a"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDBTools = "e29189f1-7114-4dbd-93d0-c5673a921a58"
 PaddedViews = "5432bcbf-9aad-5242-b902-cca2824c8663"
-Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80"
 
 [compat]
 julia = "1"

src/AssigningSecondaryStructure/dssp.jl

@@ -14,7 +14,7 @@ function _unfold(a::AbstractArray, window::Int, axis::Int)
     return _moveaxis(unfolded, axis, ndims(unfolded))
 end
 
-function get_hydrogen_positions(coord::AbstractArray{T, 3}) where T <: Real
+function _get_hydrogen_positions(coord::AbstractArray{T, 3}) where T <: Real
     vec_cn = coord[2:end, 1, :] .- coord[1:end-1, 3, :]
     vec_cn ./= mapslices(norm, vec_cn, dims=2)
     vec_can = coord[2:end, 1, :] .- coord[2:end, 2, :]
@@ -24,7 +24,7 @@ function get_hydrogen_positions(coord::AbstractArray{T, 3}) where T <: Real
     return coord[2:end, 1, :] .+ 1.01 .* vec_nh
 end
 
-function get_hbond_map(
+function _get_hbond_map(
     coord::AbstractArray{T, 3};
     cutoff::Float64 = DEFAULT_CUTOFF,
     margin::Float64 = DEFAULT_MARGIN,
@@ -36,7 +36,7 @@ function get_hbond_map(
     cpos = coord[1:end-1, 3, :]
     opos = coord[1:end-1, 4, :]
     npos = coord[2:end, 1, :]
-    hpos = get_hydrogen_positions(coord)
+    hpos = _get_hydrogen_positions(coord)
 
     cmap = repeat(reshape(cpos, 1, :, 3), outer=(residue_count-1, 1, 1))
     omap = repeat(reshape(opos, 1, :, 3), outer=(residue_count-1, 1, 1))
@@ -77,15 +77,14 @@ end
 """
     dssp(coords_chains::Vararg{AbstractArray{T, 3}, N})
 
-Takes a variable number of chains, each of which is a 3D array of shape `(residue_count, 4, 3)`.
+Takes a variable number of chains, each of which is a 3D array of shape `(3, 4, residue_count)`.
 Returns a Vector{Vector{SSClass}}, where the outer vector is the number of chains,
 and the inner vector is the secondary structure class of each residue.
 """
-function dssp(coords_chains::Vararg{AbstractArray{T, 3}, N}) where {T, N}
-    chain_lengths = size.(coords_chains, 1)
-    coords = vcat(coords_chains...)
+function dssp(coords::AbstractArray{T, 3}) where T
+    coords = permutedims(coords, (3, 2, 1))
 
-    hbmap = get_hbond_map(coords)
+    hbmap = _get_hbond_map(coords)
     hbmap = permutedims(hbmap, (2, 1))  # Rearrange to "i:C=O, j:N-H" form
 
     # Identify turn 3, 4, 5
@@ -125,6 +124,14 @@ function dssp(coords_chains::Vararg{AbstractArray{T, 3}, N}) where {T, N}
     loop = .!helix .& .!strand
 
     classes = SSClass.(findfirst.(eachrow(cat(loop, helix, strand, dims=2))))
+
+    return classes
+end
+
+function dssp(coords_chains::Vector{<:AbstractArray{T, 3}}) where T
+    chain_lengths = size.(coords_chains, 3)
+    coords = cat(coords_chains..., dims=3)
+    classes = dssp(coords)
 
     classes_chains = Vector{SSClass}[]
     i = 0
@@ -133,5 +140,4 @@ function dssp(coords_chains::Vararg{AbstractArray{T, 3}, N}) where {T, N}
     end
 
     return classes_chains
-end
-
+end

src/AssigningSecondaryStructure/pdb.jl

@@ -23,10 +23,10 @@ end
 
 function load_atom_coords(atoms::AbstractVector{<:Atom})
     residues = collect_residues(atoms)
-    coords = zeros(Float32, (length(residues), 4, 3))
+    coords = zeros(Float32, (3, 4, length(residues)))
     for (i, residue) in enumerate(residues)
         for (j, atom) in enumerate(residue)
-            coords[i, j, :] = _coords(atom)
+            coords[:, j, i] = _coords(atom)
         end
     end
     return coords
@@ -53,5 +53,5 @@ Assumes that each residue in the PDB file starts with four atoms: N, CA, C, O.
 """
 function dssp(pdb_file::String)
     coords_chains = load_pdb_coords(pdb_file)
-    return dssp(coords_chains...)
+    return dssp(coords_chains)
 end

src/BackBoner.jl

@@ -1,7 +1,5 @@
 module BackBoner
 
-export ResidueArray, Backbone, BackboneAndOxygen
-
 using LinearAlgebra
 
 include("AssigningSecondaryStructure/AssigningSecondaryStructure.jl")

src/backbone/backbone.jl

@@ -1,20 +1,43 @@
-include("utils.jl")
-include("residue_array.jl")
+"""
+    Backbone{A, T}
 
-const Backbone = ResidueArray{3}
-const BackboneAndOxygen = ResidueArray{4}
+A wrapper for a 3xAxN matrix of coordinates of atoms in a backbone chain.
+"""
+struct Backbone{A, T <: Real} <: AbstractArray{T, 3}
+    coords::AbstractArray{T, 3}
 
-nitrogen_coord_matrix(ra::ResidueArray)    = _ith_column_coords(ra, 1)
-alphacarbon_coord_matrix(ra::ResidueArray) = _ith_column_coords(ra, 2)
-carbon_coord_matrix(ra::ResidueArray)      = _ith_column_coords(ra, 3)
-oxygen_coord_matrix(ra::ResidueArray)      = _ith_column_coords(ra, 4)
+    function Backbone(coords::AbstractArray{T, 3}) where T <: Real
+        A = size(coords, 2)
+        @assert size(coords, 1) == 3 "coords must have 3 coordinates per atom"
+        return new{A, T}(coords)
+    end
+end
 
-nitrogens(ra::ResidueArray)    = eachcol(_ith_column_coords(ra, 1))
-alphacarbons(ra::ResidueArray) = eachcol(_ith_column_coords(ra, 2))
-carbons(ra::ResidueArray)      = eachcol(_ith_column_coords(ra, 3))
-oxygens(ra::ResidueArray)      = eachcol(_ith_column_coords(ra, 4))
+@inline Base.size(bb::Backbone) = size(bb.coords)
+@inline Base.length(bb::Backbone) = size(bb, 3)
+@inline Base.getindex(bb::Backbone, i, j, k) = bb.coords[i,j,k]
+@inline Base.getindex(bb::Backbone, r::UnitRange{<:Integer}) = Backbone(view(bb.coords, :, :, r))
+@inline Base.getindex(bb::Backbone, i::Integer) = view(bb.coords, :, :, i)
 
-unroll_coords(ra::ResidueArray{A}) where A = reshape(permutedims(ra.coords, (2, 1, 3)), :, A)
+"""
+    atom_coord_matrix(bb, i)
+
+Returns the coordinates of specific columns of atoms in a backbone.
+"""
+function atom_coord_matrix(bb::Backbone{A}, i) where A
+    return view(bb.coords, :, i, :)
+end
+
+"""
+    unroll_atoms(bb, i)
+
+Returns the coordinates of specific columns of atoms in a backbone,
+but unrolled into a 3xN matrix where N is the number of residues times
+the number of columns selected (atoms selected per residue).
+"""
+function unroll_atoms(bb::Backbone{A}, i=Colon()) where A
+    return reshape(atom_coord_matrix(bb, i), 3, :)
+end
 
 include("oxygen.jl")
-include("residue.jl")
+include("ncaco.jl")

src/backbone/ncaco.jl

@@ -0,0 +1,4 @@
+nitrogen_coord_matrix(bb::Backbone) = atom_coord_matrix(bb, 1)
+alphacarbon_coord_matrix(bb::Backbone) = atom_coord_matrix(bb, 2)
+carbon_coord_matrix(bb::Backbone) = atom_coord_matrix(bb, 3)
+oxygen_coord_matrix(bb::Backbone) = atom_coord_matrix(bb, 4)

src/backbone/oxygen.jl

@@ -1,3 +1,16 @@
+function get_rotation_matrix(
+    point1::AbstractVector{T}, center::AbstractVector{T}, point3::AbstractVector{T}
+) where T <: Real
+    direction1 = point3 - center
+    direction2 = point1 - center
+    basis1 = normalize(direction1)
+    orthogonal_component = direction2 - basis1 * (basis1' * direction2)
+    basis2 = normalize(orthogonal_component)
+    basis3 = cross(basis1, basis2)
+    rotation_matrix = [basis1 basis2 basis3]
+    return rotation_matrix
+end
+
 const magic_vector = [-0.672, -1.034, 0.003]
 
 function get_oxygen_coords(
@@ -8,30 +21,27 @@ function get_oxygen_coords(
 end
 
 function get_oxygen_coords(
-    backbone::Backbone{T},
+    backbone::Backbone{3, T},
 ) where T <: Real
-    L = nresidues(backbone)
+    L = length(backbone)
 
     CAs = alphacarbons(backbone)
     Cs = carbons(backbone)
     next_Ns = view(nitrogens(backbone), 2:L)
 
-    oxygen_coords = zeros(T, L-1, 3)
+    oxygen_coords = zeros(T, 3, L-1)
     for (i, (CA, C, next_N)) in enumerate(zip(CAs, Cs, next_Ns))
-        oxygen_coords[i, :] = get_oxygen_coords(CA, C, next_N)
+        oxygen_coords[:, i] = get_oxygen_coords(CA, C, next_N)
     end
 
     return oxygen_coords
 end
 
-function ResidueArray{4}(backbone::Backbone{T}) where T <: Real
-    L = nresidues(backbone)
+function backbone_with_oxygen(backbone::Backbone{3, T}) where T <: Real
+    L = length(backbone)
     oxygen_coords = get_oxygen_coords(backbone)
-    coords = zeros(T, L-1, 4, 3)
-    coords[:, 1:3, :] = backbone.coords[1:L-1, :, :]
+    coords = zeros(T, 3, 4, L-1)
+    coords[:, 1:3, :] = backbone.coords[:, :, 1:L-1]
     coords[:, 4, :] = oxygen_coords
-    return BackboneAndOxygen(coords)
-end
-
-oxygen_coord_matrix(ra::ResidueArray) = _ith_column_positions(ra, 4)
-oxygens(ra::ResidueArray) = eachrow(_ith_column_positions(ra, 4))
+    return Backbone(coords)
+end

src/visuals/render.jl

@@ -0,0 +1,45 @@
+function render!(segment::Segment{ASS.Loop})
+    startpoint = segment_startpoint(segment)
+    stoppoint = segment_stoppoint(segment)
+    anchors = alphacarbon_coord_matrix(segment.subbackbone[2:end])
+    coords = hcat(startpoint, anchors, stoppoint)
+    surface_vertices, color_matrix = tube(coords, 0.3, spline_quality=10, tube_quality=10, color_start=segment.range.start/length(segment.backbone), color_end=segment.range.stop/length(segment.backbone))
+    surface!(eachslice(surface_vertices, dims=1)..., color=color_matrix)
+end
+
+function render!(segment::Segment{ASS.Helix})
+    startpoint = segment_startpoint(segment)
+    stoppoint = segment_stoppoint(segment)
+    anchors = alphacarbon_coord_matrix(segment.subbackbone[2:end]) # startpoint is first point instead. including first N *and* CA could mess with normals
+    coords = hcat(startpoint, anchors, stoppoint)
+    surface_vertices, color_matrix = tube(coords, 1, x_elongation=0.15, color_start=segment.range.start/length(segment.backbone), color_end=segment.range.stop/length(segment.backbone))
+    surface!(eachslice(surface_vertices, dims=1)..., color=color_matrix)
+end
+
+function render!(segment::Segment{ASS.Strand})
+    startpoint = segment_startpoint(segment)
+    stoppoint = segment_stoppoint(segment)
+    oxygen_coords = oxygen_coord_matrix(segment.subbackbone)
+    oxygen_coords_side1 = oxygen_coords[:, 1:2:end-1]
+    oxygen_coords_side2 = oxygen_coords[:, 2:2:end]
+    coords1 = hcat(startpoint, oxygen_coords_side1, stoppoint)
+    coords2 = hcat(startpoint, oxygen_coords_side2, stoppoint)
+    surface_vertices, color_matrix = sheet(coords1, coords2, thickness=0.5, width_pad=0.5, color_start=segment.range.start/length(segment.backbone), color_end=segment.range.stop/length(segment.backbone))
+    surface!(eachslice(surface_vertices, dims=1)..., color=color_matrix)
+end
+
+function render!(backbone::Backbone{4}, ssclasses::Vector{ASS.SSClass})
+    segments = segmenter(ssclasses, backbone)
+    render!.(segments)
+end
+
+function render!(backbones::Vector{<:Backbone{4}})
+    set_theme!(backgroundcolor = :black)
+
+    ssclasses_vec = ASS.dssp([bb.coords for bb in backbones])
+    scene = lines(Float64[])#, axis=(;type=Axis3, aspect=:data))
+    for (backbone, ssclasses) in zip(backbones, ssclasses_vec)
+        render!(backbone, ssclasses)
+    end
+    display(scene)
+end

src/visuals/segment.jl

@@ -0,0 +1,37 @@
+struct Segment{C}
+    range::UnitRange{Int}
+    backbone::Backbone
+    subbackbone::Backbone
+
+    function Segment{C}(range::UnitRange{Int}, backbone::Backbone) where C
+        return new{C}(range, backbone, backbone[range])
+    end
+end
+
+function segmenter(ss::AbstractVector{ASS.SSClass}, backbone::Backbone)
+    start_idx = 1
+    end_idx = 1
+    segments = Segment[]
+    for (i, class) in enumerate(ss)
+        if class != ss[start_idx]
+            push!(segments, Segment{ss[start_idx]}(start_idx:end_idx, backbone))
+            start_idx = i
+        end
+        end_idx = i
+    end
+    push!(segments, Segment{ss[start_idx]}(start_idx:end_idx, backbone))
+    return segments
+end
+
+function segment_startpoint(segment::Segment)
+    return segment.subbackbone[1][:,1] # first nitrogen atom
+end
+
+function segment_stoppoint(segment::Segment)
+    stop = segment.range.stop
+    if stop == length(segment.backbone)
+        return segment.backbone[stop][:,3] # last carbon atom
+    else
+        return segment.backbone[stop+1][:,1] # next nitrogen atom
+    end
+end