345 lines
12 KiB
Python
345 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import math
|
|
from dataclasses import dataclass, field
|
|
from html import escape
|
|
|
|
|
|
NODE_RADIUS = 10.0
|
|
PRIMARY_SPACE = 20.0
|
|
PAIR_SPACE = 20.0
|
|
PADDING = 52.0
|
|
|
|
|
|
def _fmt(value):
|
|
return f"{float(value):.2f}"
|
|
|
|
|
|
def _fmt3(value):
|
|
return f"{float(value):.3f}"
|
|
|
|
|
|
def _pairmap_from_structure(structure):
|
|
pair_stack = []
|
|
dangling_stack = []
|
|
pairs = [-1] * len(structure)
|
|
|
|
for index, char in enumerate(structure):
|
|
if char == "(":
|
|
pair_stack.append(index)
|
|
elif char == ")":
|
|
if pair_stack:
|
|
partner = pair_stack.pop()
|
|
pairs[index] = partner
|
|
pairs[partner] = index
|
|
else:
|
|
dangling_stack.append(index)
|
|
|
|
if pair_stack:
|
|
if len(pair_stack) != len(dangling_stack):
|
|
raise ValueError("Unbalanced structure for split-strand layout.")
|
|
for left, right in zip(pair_stack, reversed(dangling_stack)):
|
|
pairs[left] = right
|
|
pairs[right] = left
|
|
|
|
return pairs
|
|
|
|
|
|
@dataclass
|
|
class _TreeNode:
|
|
children: list["_TreeNode"] = field(default_factory=list)
|
|
is_pair: bool = False
|
|
index_a: int = -1
|
|
index_b: int = -1
|
|
x: float = 0.0
|
|
y: float = 0.0
|
|
go_x: float = 0.0
|
|
go_y: float = 0.0
|
|
|
|
|
|
def _add_nodes_recursive(pairmap, root, start, end):
|
|
if start > end:
|
|
raise ValueError(f"Invalid recursive span for split-strand layout: {start}>{end}")
|
|
|
|
if pairmap[start] == end:
|
|
child = _TreeNode(is_pair=True, index_a=start, index_b=end)
|
|
_add_nodes_recursive(pairmap, child, start + 1, end - 1)
|
|
root.children.append(child)
|
|
return
|
|
|
|
child = _TreeNode()
|
|
cursor = start
|
|
while cursor <= end:
|
|
partner = pairmap[cursor]
|
|
if partner > cursor:
|
|
_add_nodes_recursive(pairmap, child, cursor, partner)
|
|
cursor = partner + 1
|
|
continue
|
|
child.children.append(_TreeNode(index_a=cursor))
|
|
cursor += 1
|
|
root.children.append(child)
|
|
|
|
|
|
def _setup_coords_recursive(node, parent, start_x, start_y, go_x, go_y, flipped=False):
|
|
cross_x = -go_y
|
|
cross_y = go_x
|
|
node.go_x = go_x
|
|
node.go_y = go_y
|
|
|
|
if len(node.children) == 1:
|
|
node.x = start_x
|
|
node.y = start_y
|
|
child = node.children[0]
|
|
next_x = start_x + go_x * PRIMARY_SPACE
|
|
next_y = start_y + (-1 if flipped else 1) * go_y * PRIMARY_SPACE
|
|
if child.is_pair or (not child.is_pair and child.index_a >= 0):
|
|
_setup_coords_recursive(child, node, next_x, next_y, go_x, go_y, flipped=flipped)
|
|
else:
|
|
_setup_coords_recursive(child, node, start_x, start_y, go_x, go_y, flipped=flipped)
|
|
return
|
|
|
|
if not node.children:
|
|
node.x = start_x
|
|
node.y = start_y
|
|
return
|
|
|
|
pair_count = sum(1 for child in node.children if child.is_pair)
|
|
circle_length = (len(node.children) + 1) * PRIMARY_SPACE + (pair_count + 1) * PAIR_SPACE
|
|
circle_radius = circle_length / (2 * math.pi)
|
|
length_walker = PAIR_SPACE / 2.0
|
|
|
|
if parent is None:
|
|
node.x = go_x * circle_radius
|
|
node.y = go_y * circle_radius
|
|
else:
|
|
node.x = parent.x + go_x * circle_radius
|
|
node.y = parent.y + (-1 if flipped else 1) * go_y * circle_radius
|
|
|
|
for child in node.children:
|
|
length_walker += PRIMARY_SPACE
|
|
if child.is_pair:
|
|
length_walker += PAIR_SPACE / 2.0
|
|
|
|
rad_angle = length_walker / circle_length * 2 * math.pi - math.pi / 2.0
|
|
if parent is None:
|
|
rad_angle -= math.pi / 2.0
|
|
|
|
child_x = node.x + math.cos(rad_angle) * cross_x * circle_radius + math.sin(rad_angle) * go_x * circle_radius
|
|
child_y = node.y + (-1 if flipped else 1) * math.cos(rad_angle) * cross_y * circle_radius + (-1 if flipped else 1) * math.sin(rad_angle) * go_y * circle_radius
|
|
|
|
child_go_x = child_x - node.x
|
|
child_go_y = child_y - node.y
|
|
child_go_len = math.hypot(child_go_x, child_go_y) or 1.0
|
|
|
|
_setup_coords_recursive(
|
|
child,
|
|
node,
|
|
child_x,
|
|
child_y,
|
|
child_go_x / child_go_len,
|
|
(-1 if flipped else 1) * child_go_y / child_go_len,
|
|
flipped=flipped,
|
|
)
|
|
|
|
if child.is_pair:
|
|
length_walker += PAIR_SPACE / 2.0
|
|
|
|
|
|
def _collect_coords_recursive(node, xs, ys, flipped=False):
|
|
if node.is_pair:
|
|
cross_x = -node.go_y
|
|
cross_y = node.go_x
|
|
xs[node.index_a] = node.x + cross_x * PAIR_SPACE / 2.0
|
|
xs[node.index_b] = node.x - cross_x * PAIR_SPACE / 2.0
|
|
ys[node.index_a] = node.y + (-1 if flipped else 1) * cross_y * PAIR_SPACE / 2.0
|
|
ys[node.index_b] = node.y + (1 if flipped else -1) * cross_y * PAIR_SPACE / 2.0
|
|
elif node.index_a >= 0:
|
|
xs[node.index_a] = node.x
|
|
ys[node.index_a] = node.y
|
|
|
|
for child in node.children:
|
|
_collect_coords_recursive(child, xs, ys, flipped=flipped)
|
|
|
|
|
|
def _layout_positions(display_structure):
|
|
pairmap = _pairmap_from_structure(display_structure)
|
|
root = _TreeNode()
|
|
cursor = 0
|
|
while cursor < len(pairmap):
|
|
partner = pairmap[cursor]
|
|
if partner > cursor:
|
|
_add_nodes_recursive(pairmap, root, cursor, partner)
|
|
cursor = partner + 1
|
|
continue
|
|
root.children.append(_TreeNode(index_a=cursor))
|
|
cursor += 1
|
|
|
|
xs = [0.0] * len(display_structure)
|
|
ys = [0.0] * len(display_structure)
|
|
_setup_coords_recursive(root, None, 0.0, 0.0, 0.0, 1.0, flipped=False)
|
|
_collect_coords_recursive(root, xs, ys, flipped=False)
|
|
|
|
min_x = min(x - NODE_RADIUS for x in xs)
|
|
min_y = min(y - NODE_RADIUS for y in ys)
|
|
xs = [x - min_x for x in xs]
|
|
ys = [y - min_y for y in ys]
|
|
return pairmap, xs, ys
|
|
|
|
|
|
def _strand_spans(display_sequence):
|
|
spans = []
|
|
start = None
|
|
for index, char in enumerate(display_sequence):
|
|
if char == " ":
|
|
if start is not None:
|
|
spans.append((start, index - 1))
|
|
start = None
|
|
continue
|
|
if start is None:
|
|
start = index
|
|
if start is not None:
|
|
spans.append((start, len(display_sequence) - 1))
|
|
return spans
|
|
|
|
|
|
def render_split_strands_svg(strand_sequences, structure, title=None):
|
|
if not strand_sequences:
|
|
raise ValueError("Split-strand layout requires at least one strand.")
|
|
|
|
display_sequence = " ".join(str(sequence) for sequence in strand_sequences)
|
|
display_structure = str(structure).replace("+", " ")
|
|
|
|
if len(display_sequence) != len(display_structure):
|
|
raise ValueError("Sequence and structure length mismatch for split-strand layout.")
|
|
|
|
pairmap, xs, ys = _layout_positions(display_structure)
|
|
visible_indices = [index for index, char in enumerate(display_sequence) if char != " "]
|
|
if not visible_indices:
|
|
raise ValueError("Split-strand layout produced no visible residues.")
|
|
|
|
min_x = min(xs[index] for index in visible_indices) - NODE_RADIUS - PADDING
|
|
min_y = min(ys[index] for index in visible_indices) - NODE_RADIUS - PADDING
|
|
max_x = max(xs[index] for index in visible_indices) + NODE_RADIUS + PADDING
|
|
max_y = max(ys[index] for index in visible_indices) + NODE_RADIUS + PADDING
|
|
|
|
width = max_x - min_x
|
|
height = max_y - min_y
|
|
|
|
def tx(index):
|
|
return xs[index] - min_x
|
|
|
|
def ty(index):
|
|
return ys[index] - min_y
|
|
|
|
compact_index = {}
|
|
compact_counter = 1
|
|
for index in visible_indices:
|
|
compact_index[index] = compact_counter
|
|
compact_counter += 1
|
|
|
|
svg_lines = [
|
|
'<?xml version="1.0" encoding="UTF-8" standalone="yes" ?>',
|
|
(
|
|
f'<svg xmlns="http://www.w3.org/2000/svg" width="{_fmt(width)}" height="{_fmt(height)}" '
|
|
f'viewBox="0 0 {_fmt(width)} {_fmt(height)}" preserveAspectRatio="xMidYMid meet" '
|
|
'data-plot-engine="split-strands">'
|
|
),
|
|
]
|
|
if title:
|
|
svg_lines.append(f"<title>{escape(title)}</title>")
|
|
svg_lines.extend(
|
|
[
|
|
' <script type="text/ecmascript">',
|
|
" <![CDATA[",
|
|
" var shown = 1;",
|
|
" function click() {",
|
|
' var seq = document.getElementById("seq");',
|
|
" if (shown==1) {",
|
|
' seq.setAttribute("style", "visibility: hidden");',
|
|
" shown = 0;",
|
|
" } else {",
|
|
' seq.setAttribute("style", "visibility: visible");',
|
|
" shown = 1;",
|
|
" }",
|
|
" }",
|
|
" ]]>",
|
|
" </script>",
|
|
' <style type="text/css">',
|
|
" <![CDATA[",
|
|
" .nucleotide {",
|
|
" font-family: SansSerif;",
|
|
" }",
|
|
" .backbone {",
|
|
" stroke: grey;",
|
|
" fill: none;",
|
|
" stroke-width: 1.5;",
|
|
" }",
|
|
" .basepairs {",
|
|
" stroke: red;",
|
|
" fill: none;",
|
|
" stroke-width: 2.5;",
|
|
" }",
|
|
" ]]>",
|
|
" </style>",
|
|
"",
|
|
f' <rect style="stroke: white; fill: white" height="{_fmt(height)}" x="0" y="0" width="{_fmt(width)}" onclick="click(evt)" />',
|
|
' <g transform="scale(1,1) translate(0,0)">',
|
|
]
|
|
)
|
|
|
|
for strand_index, (start, end) in enumerate(_strand_spans(display_sequence), start=1):
|
|
points = " ".join(
|
|
f"{_fmt3(tx(index))},{_fmt3(ty(index))}"
|
|
for index in range(start, end + 1)
|
|
)
|
|
svg_lines.append(f' <polyline class="backbone" id="outline-{strand_index}" points=" {points} " />')
|
|
|
|
svg_lines.append(' <g id="pairs">')
|
|
for left, right in enumerate(pairmap):
|
|
if right <= left or display_sequence[left] == " " or display_sequence[right] == " ":
|
|
continue
|
|
svg_lines.append(
|
|
f' <line class="basepairs" id="{compact_index[left]},{compact_index[right]}" '
|
|
f'x1="{_fmt(tx(left))}" y1="{_fmt(ty(left))}" '
|
|
f'x2="{_fmt(tx(right))}" y2="{_fmt(ty(right))}" />'
|
|
)
|
|
svg_lines.append(" </g>")
|
|
|
|
svg_lines.append(' <g transform="translate(-4.6, 4)" id="seq">')
|
|
for index in visible_indices:
|
|
base = display_sequence[index]
|
|
svg_lines.append(
|
|
f' <text class="nucleotide" x="{_fmt3(tx(index))}" y="{_fmt3(ty(index))}">{escape(base)}</text>'
|
|
)
|
|
svg_lines.append(" </g>")
|
|
|
|
compact_sequence = display_sequence.replace(" ", "")
|
|
basepair_rows = [
|
|
f' {{ i: {compact_index[left]}, j: {compact_index[right]}, type: "cWW" }}'
|
|
for left, right in enumerate(pairmap)
|
|
if right > left and display_sequence[left] != " " and display_sequence[right] != " "
|
|
]
|
|
coord_rows = [
|
|
f' {{ x: {_fmt3(tx(index))}, y: {_fmt3(ty(index))} }}'
|
|
for index in visible_indices
|
|
]
|
|
svg_lines.extend(
|
|
[
|
|
" </g>",
|
|
'<script type="text/ecmascript">',
|
|
"<![CDATA[",
|
|
f" let sequence = {json.dumps(compact_sequence)};",
|
|
f" let structure = {json.dumps(str(structure))};",
|
|
" const basepairs = [",
|
|
",\n".join(basepair_rows) + ("" if not basepair_rows else ""),
|
|
" ];",
|
|
" const coords = [",
|
|
",\n".join(coord_rows) + ("" if not coord_rows else ""),
|
|
" ];",
|
|
"]]>",
|
|
"</script>",
|
|
"</svg>",
|
|
]
|
|
)
|
|
return "\n".join(svg_lines)
|