From de673e8f2330c29b8a342d585e5511a39b9435f7 Mon Sep 17 00:00:00 2001
From: anton083 <anton.oresten42@gmail.com>
Date: Fri, 5 Apr 2024 12:49:06 +0200
Subject: [PATCH] Draw line segments when there are missing residues

---
 src/ribbon/render.jl | 18 +++++++++++++++++-
 src/ribbon/utils.jl  | 21 +++++++++++++++++++++
 2 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/src/ribbon/render.jl b/src/ribbon/render.jl
index 34ef3ee..39ac930 100644
--- a/src/ribbon/render.jl
+++ b/src/ribbon/render.jl
@@ -103,13 +103,29 @@ function render!(
     protein::AbstractVector{Protein.Chain};
     colorscheme::ColorScheme = default_colorscheme,
     color_vectors::AbstractVector{<:AbstractVector{<:RGB}} = [colorscheme[LinRange(0, 1, length(chain))] for chain in protein],
+    plots = nothing,
     kwargs...
 )
     @assert Protein.has_assigned_ss(protein) "Protein must have assigned secondary structure."
     @assert length(protein) == length(color_vectors)
     @assert length.(protein) == length.(color_vectors)
     for (chain, colors) in zip(protein, color_vectors)
-        render!(container, chain, colors; kwargs...)
+        subchain_ranges = split_by_resnum(chain)
+        if length(subchain_ranges) == 1
+            render!(container, chain, colors; kwargs...)
+        else
+            subchains = [chain[r] for r in subchain_ranges]
+            for (subchain, r) in zip(subchains, subchain_ranges)
+                render!(container, subchain, colors[r]; kwargs...)
+            end
+            # draw lines between start and ends of subchains
+            for (i, j) in zip(1:length(subchains)-1, 2:length(subchains))
+                startpoint, endpoint = subchains[i].backbone[end], subchains[j].backbone[begin]
+                xs, ys, zs = [LinRange(startpoint[i], endpoint[i], 16) for i in 1:3]
+                p = linesegments!(container, xs, ys, zs; linewidth=3, color=colorant"white", transparency=true)
+                !isnothing(plots) && push!(plots, p)
+            end
+        end
     end
 
     return container
diff --git a/src/ribbon/utils.jl b/src/ribbon/utils.jl
index a57a7ff..79a271f 100644
--- a/src/ribbon/utils.jl
+++ b/src/ribbon/utils.jl
@@ -18,4 +18,25 @@ function expand_colors(colors::Vector, N::Integer)
     end
 
     return reshape(result[1:N], :, 1)
+end
+
+function Base.getindex(chain::Protein.Chain, r::UnitRange{<:Integer})
+    Protein.Chain(chain.id, chain.backbone[3*r.start-2:3r.stop], modelnum=chain.modelnum, resnums=chain.resnums[r], aavector=chain.aavector[r], ssvector=chain.ssvector[r])
+end
+
+# split a chain into subchains if there is a gap in the residue numbering
+function split_by_resnum(chain::Protein.Chain)
+    split_indices = findall(diff(chain.resnums) .!= 1)
+    if isempty(split_indices)
+        return [chain]
+    end
+
+    ranges = UnitRange{Int}[]
+    start_idx = 0
+    for idx in split_indices
+        push!(ranges, start_idx+1:idx)
+        start_idx = idx
+    end
+    push!(ranges, start_idx+1:length(chain))
+    return ranges
 end
\ No newline at end of file