use std::sync::Arc;
use rstui::{App, AppContext, AppEvent, ControlFlow, Frame, Rgba, TextRun};
use crate::atlas::AtlasConfidence;
use crate::pdb::{Atom, Protein, SecondaryStructure, Vec3};
use crate::raster::{Canvas, Color, mix, scale};
const BACKGROUND_TOP: Color = [5, 9, 17, 255];
const BACKGROUND_BOTTOM: Color = [12, 18, 29, 255];
const TEXT: Color = [220, 230, 241, 255];
const MUTED: Color = [135, 153, 174, 255];
const ACCENT: Color = [73, 215, 196, 255];
const PANEL_BACKGROUND: Color = [13, 6, 23, 255];
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum Representation {
Backbone,
Atoms,
Combined,
}
impl Representation {
pub const fn name(self) -> &'static str {
match self {
Self::Backbone => "ribbon",
Self::Atoms => "atoms",
Self::Combined => "combined",
}
}
fn next(self) -> Self {
match self {
Self::Backbone => Self::Atoms,
Self::Atoms => Self::Combined,
Self::Combined => Self::Backbone,
}
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ColorScheme {
Chain,
Element,
Confidence,
}
impl ColorScheme {
pub const fn name(self) -> &'static str {
match self {
Self::Chain => "chain",
Self::Element => "element",
Self::Confidence => "pLDDT/B-factor",
}
}
fn next(self) -> Self {
match self {
Self::Chain => Self::Element,
Self::Element => Self::Confidence,
Self::Confidence => Self::Chain,
}
}
const fn compact_name(self) -> &'static str {
match self {
Self::Chain => "chain",
Self::Element => "element",
Self::Confidence => "pLDDT",
}
}
}
pub struct Viewer {
protein: Arc<Protein>,
view_center: Vec3,
yaw: f32,
pitch: f32,
zoom: f32,
representation: Representation,
color_scheme: ColorScheme,
dragging: Option<(u32, u32)>,
click_candidate: Option<usize>,
dragged: bool,
hovered_residue: Option<usize>,
selected_residue: Option<usize>,
auto_rotate: bool,
last_tick_ms: u64,
atlas: Option<AtlasPanel>,
}
struct AtlasPanel {
label: String,
confidence: Option<AtlasConfidence>,
}
impl Viewer {
pub fn new(
protein: Protein,
representation: Representation,
color_scheme: ColorScheme,
) -> Self {
let view_center = protein.center;
Self {
protein: Arc::new(protein),
view_center,
yaw: -0.72,
pitch: 0.38,
zoom: 1.35,
representation,
color_scheme,
dragging: None,
click_candidate: None,
dragged: false,
hovered_residue: None,
selected_residue: None,
auto_rotate: false,
last_tick_ms: 0,
atlas: None,
}
}
pub fn with_atlas(
mut self,
label: impl Into<String>,
confidence: Option<AtlasConfidence>,
) -> Self {
self.atlas = Some(AtlasPanel {
label: label.into(),
confidence,
});
self
}
pub fn with_selected_residue(mut self, chain: &str, number: i32) -> Result<Self, String> {
let index = self
.protein
.residues
.iter()
.position(|residue| residue.chain == chain && residue.number == number)
.ok_or_else(|| format!("residue {chain}:{number} was not found"))?;
self.hovered_residue = Some(index);
self.set_selection(Some(index));
Ok(self)
}
fn reset_view(&mut self) {
self.view_center = self.protein.center;
self.yaw = -0.72;
self.pitch = 0.38;
self.zoom = 1.35;
self.selected_residue = None;
self.auto_rotate = false;
}
pub fn snapshot(&self, width: u32, height: u32) -> Frame {
self.render_scene(width.max(1), height.max(1), false, None)
}
fn render_scene(
&self,
width: u32,
height: u32,
fast: bool,
terminal: Option<(u16, u16)>,
) -> Frame {
let mut canvas = Canvas::new(width, height, BACKGROUND_TOP);
canvas.vertical_gradient(BACKGROUND_TOP, BACKGROUND_BOTTOM);
let layout = terminal.map_or_else(
|| self.view_layout(width, height),
|(columns, rows)| self.terminal_view_layout(width, height, columns, rows),
);
let native_text = terminal.is_some();
canvas.set_text_scale(layout.font_scale);
let scene_center = layout.scene_center;
let scene_scale = layout.scene_scale;
let projected: Vec<Projected> = self
.protein
.atoms
.iter()
.map(|atom| self.project_position(atom.position, scene_center, scene_scale))
.collect();
if matches!(
self.representation,
Representation::Backbone | Representation::Combined
) {
self.render_ribbon(&mut canvas, scene_center, scene_scale, fast);
}
if matches!(
self.representation,
Representation::Atoms | Representation::Combined
) {
let maximum = if fast { 8_000 } else { 30_000 };
let stride = self.protein.atoms.len().div_ceil(maximum).max(1);
for (index, atom) in self.protein.atoms.iter().enumerate().step_by(stride) {
let point = projected[index];
let physical_radius = element_radius(&atom.element);
let radius = (physical_radius / self.protein.radius * scene_scale)
.clamp(if fast { 1.0 } else { 1.5 }, if fast { 7.0 } else { 12.0 });
let color = self.atom_color(atom, point.z);
canvas.sphere(point.x, point.y, point.z + 0.002, radius, color);
}
}
if let Some(selected) = self.selected_residue {
self.render_selection_detail(
&mut canvas,
selected,
&projected,
scene_center,
scene_scale,
);
}
if let Some(atlas) = &self.atlas
&& layout.molecule_width < width
{
let panel = Rect {
x: layout.molecule_width,
y: layout.scene_top,
width: width - layout.molecule_width,
height: layout.scene_height,
};
let ui_cell_width = terminal.map_or_else(
|| canvas.text_cell_width(),
|(columns, _)| (width / u32::from(columns.max(1))).max(1),
);
let ui_cell_height = terminal.map_or_else(
|| canvas.text_cell_height(),
|(_, rows)| (height / u32::from(rows.max(1))).max(1),
);
render_pae_panel(
&mut canvas,
panel,
atlas,
!native_text,
ui_cell_width,
ui_cell_height,
);
}
if layout.detailed {
canvas.fill_rect(0, 0, width, layout.title_height, [3, 7, 13, 255]);
canvas.fill_rect(
0,
height.saturating_sub(layout.footer_height),
width,
layout.footer_height,
[3, 7, 13, 255],
);
canvas.horizontal_line(layout.title_height.saturating_sub(1), [35, 55, 71, 255]);
canvas.horizontal_line(
height.saturating_sub(layout.footer_height + 1),
[35, 55, 71, 255],
);
let cell_width = canvas.text_cell_width();
let cell_height = canvas.text_cell_height();
let title_capacity =
layout.molecule_width.saturating_sub(16) as usize / cell_width as usize;
if !native_text {
canvas.text(
8,
layout.title_height.saturating_sub(cell_height) / 2,
&truncate(&self.protein.title, title_capacity),
TEXT,
);
}
if layout.sequence_height > 0 {
self.render_sequence_panel(&mut canvas, layout, !native_text);
}
if let Some(residue) = self.selected_residue.or(self.hovered_residue) {
self.render_residue_card(&mut canvas, residue, layout, !native_text, terminal);
}
let status = format!(
"{} atoms | {} residues | {} {} | {} | {}{}",
self.protein.atoms.len(),
self.protein.residue_count,
self.protein.chains.len(),
if self.protein.chains.len() == 1 {
"chain"
} else {
"chains"
},
self.representation.name(),
self.color_scheme.compact_name(),
if self.auto_rotate { " | AUTO" } else { "" }
);
let footer_y = height
.saturating_sub(layout.footer_height)
.saturating_add(layout.footer_height.saturating_sub(cell_height) / 2);
if !native_text {
canvas.text(8, footer_y, &status, ACCENT);
}
let help = "hover inspect click detail drag rotate wheel zoom m mode c color x clear r reset q quit";
let help_x = width.saturating_sub(help.len() as u32 * cell_width + 8);
if !native_text && help_x > status.len() as u32 * cell_width + cell_width * 3 {
canvas.text(help_x, footer_y, help, MUTED);
}
}
let mut frame =
Frame::rgba(width, height, canvas.pixels).expect("canvas dimensions are valid");
if let Some((columns, rows)) = terminal {
self.add_terminal_text(&mut frame, layout, width, height, columns, rows);
}
frame
}
fn view_layout(&self, width: u32, height: u32) -> ViewLayout {
let detailed = width >= 280 && height >= 105;
let font_scale = if width >= 700 && height >= 400 { 2 } else { 1 };
let glyph_height = 8 * font_scale;
let title_height = if detailed { glyph_height + 12 } else { 2 };
let sequence_height = if detailed && height >= 220 {
if font_scale == 2 { 106 } else { 70 }
} else {
0
};
let footer_height = if detailed { glyph_height + 12 } else { 2 };
let scene_top = title_height + sequence_height;
let scene_height = height.saturating_sub(scene_top + footer_height).max(1);
let molecule_width = if self.atlas.is_some() && width >= 520 {
(width as f32 * 0.62).round() as u32
} else {
width
}
.max(1);
let scene_center = [
molecule_width as f32 * 0.5,
scene_top as f32 + scene_height as f32 * 0.5,
];
let scene_scale = molecule_width.min(scene_height) as f32 * 0.43 * self.zoom;
ViewLayout {
detailed,
font_scale,
title_height,
footer_height,
sequence_height,
scene_top,
scene_height,
molecule_width,
scene_center,
scene_scale,
}
}
fn terminal_view_layout(&self, width: u32, height: u32, columns: u16, rows: u16) -> ViewLayout {
let detailed = columns >= 60 && rows >= 18;
let row_height = height as f32 / u32::from(rows.max(1)) as f32;
let rows_to_pixels =
|count: u32| ((row_height * count as f32).round() as u32).max(count.min(height));
let title_height = if detailed { rows_to_pixels(2) } else { 2 };
let sequence_height = if detailed && rows >= 22 {
rows_to_pixels(5)
} else {
0
};
let footer_height = if detailed { rows_to_pixels(2) } else { 2 };
let scene_top = title_height + sequence_height;
let scene_height = height.saturating_sub(scene_top + footer_height).max(1);
let molecule_width = if self.atlas.is_some() && columns >= 80 {
(width as f32 * 0.62).round() as u32
} else {
width
}
.max(1);
let scene_center = [
molecule_width as f32 * 0.5,
scene_top as f32 + scene_height as f32 * 0.5,
];
let scene_scale = molecule_width.min(scene_height) as f32 * 0.43 * self.zoom;
ViewLayout {
detailed,
font_scale: 1,
title_height,
footer_height,
sequence_height,
scene_top,
scene_height,
molecule_width,
scene_center,
scene_scale,
}
}
fn render_sequence_panel(&self, canvas: &mut Canvas, layout: ViewLayout, draw_text: bool) {
let panel_y = layout.title_height;
let cell_width = canvas.text_cell_width();
let cell_height = canvas.text_cell_height();
canvas.fill_rect(
0,
panel_y,
canvas.width,
layout.sequence_height,
[8, 12, 20, 255],
);
canvas.horizontal_line(layout.scene_top.saturating_sub(1), [49, 60, 78, 255]);
if !draw_text {
return;
}
let active = self.selected_residue.or(self.hovered_residue);
let chain = active
.map(|index| self.protein.residues[index].chain.as_str())
.or_else(|| self.protein.chains.first().map(String::as_str))
.unwrap_or("_");
let residue_indices: Vec<usize> = self
.protein
.residues
.iter()
.enumerate()
.filter_map(|(index, residue)| {
(residue.chain == chain && residue.ca.is_some()).then_some(index)
})
.collect();
if residue_indices.is_empty() {
return;
}
let header = format!(
"SEQUENCE CHAIN {chain} {} RESIDUES hover inspect / click select",
residue_indices.len()
);
canvas.text(
10,
panel_y + 4,
&truncate(&header, canvas.width as usize / cell_width as usize),
TEXT,
);
let per_row = (canvas.width.saturating_sub(24) / cell_width).max(10) as usize;
let capacity = per_row * 2;
let active_position = active.and_then(|active| {
residue_indices
.iter()
.position(|residue| *residue == active)
});
let mut start = active_position
.map(|position| position.saturating_sub(capacity / 2))
.unwrap_or(0)
.min(residue_indices.len().saturating_sub(capacity));
start = start / 10 * 10;
for visible in 0..capacity.min(residue_indices.len().saturating_sub(start)) {
let sequence_position = start + visible;
let row = visible / per_row;
let column = visible % per_row;
let residue_index = residue_indices[sequence_position];
let residue = &self.protein.residues[residue_index];
let x = 12 + column as u32 * cell_width;
let number_y = panel_y + cell_height + 8 + row as u32 * (cell_height * 2 + 8);
let residue_y = number_y + cell_height + 2;
if sequence_position.is_multiple_of(10) {
canvas.text(
x,
number_y,
&residue.number.to_string(),
[75, 164, 255, 255],
);
}
if self.selected_residue == Some(residue_index) {
canvas.fill_rect(
x,
residue_y.saturating_sub(1),
cell_width,
cell_height + 2,
[190, 35, 111, 255],
);
} else if self.hovered_residue == Some(residue_index) {
canvas.fill_rect(
x,
residue_y.saturating_sub(1),
cell_width,
cell_height + 2,
[112, 42, 90, 255],
);
}
let letter = amino_acid_letter(&residue.name);
canvas.text(x, residue_y, &letter.to_string(), TEXT);
}
}
fn render_residue_card(
&self,
canvas: &mut Canvas,
residue_index: usize,
layout: ViewLayout,
draw_text: bool,
terminal: Option<(u16, u16)>,
) {
let residue = &self.protein.residues[residue_index];
let atoms = self.residue_atom_indices(residue_index);
let confidence = residue
.ca
.map(|index| self.protein.confidence(self.protein.atoms[index].b_factor));
let selected = self.selected_residue == Some(residue_index);
let cell_width = canvas.text_cell_width();
let cell_height = canvas.text_cell_height();
let card = if let Some((columns, rows)) = terminal {
terminal_card_layout(
layout,
TerminalSurface {
width: canvas.width,
height: canvas.height,
columns,
rows,
},
)
} else {
let target_width = if layout.font_scale == 2 { 780 } else { 400 };
let width = layout.molecule_width.saturating_sub(20).min(target_width);
let height = (cell_height * 2 + 22).max(4);
TerminalCardLayout {
column: 0,
row: 0,
width_columns: 0,
rect: Rect {
x: layout.molecule_width.saturating_sub(width + 10),
y: layout
.scene_top
.saturating_add(layout.scene_height)
.saturating_sub(height + 10),
width,
height,
},
}
};
canvas.fill_rect(
card.rect.x,
card.rect.y,
card.rect.width,
card.rect.height,
[3, 7, 13, 232],
);
canvas.stroke_rect(
card.rect.x,
card.rect.y,
card.rect.width,
card.rect.height,
if selected {
[247, 66, 151, 255]
} else {
[78, 103, 132, 255]
},
);
if !draw_text {
return;
}
let insertion = if residue.insertion_code == ' ' {
String::new()
} else {
residue.insertion_code.to_string()
};
let first = format!(
"{} {} | {} {}{} {} {} atoms",
if selected { "SELECTED" } else { "HOVER" },
residue.chain,
residue.name,
residue.number,
insertion,
residue.secondary.name(),
atoms.len()
);
let capacity = card.rect.width as usize / cell_width as usize;
canvas.text(
card.rect.x + 8,
card.rect.y + 7,
&truncate(&first, capacity.saturating_sub(2)),
TEXT,
);
let second = if let Some(confidence) = confidence {
format!(
"pLDDT {:.2} ({}){}",
confidence,
confidence_label(confidence),
if selected {
" click again or x to clear"
} else {
" click for atom detail"
}
)
} else if selected {
"click again or x to clear".to_owned()
} else {
"click for atom detail".to_owned()
};
canvas.text(
card.rect.x + 8,
card.rect.y + 11 + cell_height,
&truncate(&second, capacity.saturating_sub(2)),
ACCENT,
);
}
fn add_terminal_text(
&self,
frame: &mut Frame,
layout: ViewLayout,
width: u32,
height: u32,
columns: u16,
rows: u16,
) {
if !layout.detailed || columns == 0 || rows == 0 {
return;
}
let title_background = [3, 7, 13, 255];
let sequence_background = [8, 12, 20, 255];
let card_background = [3, 7, 13, 255];
let surface = TerminalSurface {
width,
height,
columns,
rows,
};
let molecule_columns = pixel_to_column_edge(layout.molecule_width, width, columns).max(1);
push_terminal_text(
frame,
1,
0,
truncate(
&self.protein.title,
usize::from(molecule_columns.saturating_sub(2)),
),
TEXT,
title_background,
true,
);
if layout.sequence_height > 0 {
self.add_terminal_sequence(
frame,
columns,
pixel_to_row(layout.title_height, height, rows),
sequence_background,
);
}
if let Some(residue_index) = self.selected_residue.or(self.hovered_residue) {
let residue = &self.protein.residues[residue_index];
let atoms = self.residue_atom_indices(residue_index);
let selected = self.selected_residue == Some(residue_index);
let insertion = if residue.insertion_code == ' ' {
String::new()
} else {
residue.insertion_code.to_string()
};
let card = terminal_card_layout(layout, surface);
let first = format!(
"{} {} | {} {}{} {} {} atoms",
if selected { "SELECTED" } else { "HOVER" },
residue.chain,
residue.name,
residue.number,
insertion,
residue.secondary.name(),
atoms.len()
);
push_terminal_text(
frame,
card.column + 1,
card.row,
truncate(&first, usize::from(card.width_columns.saturating_sub(2))),
TEXT,
card_background,
true,
);
let confidence = residue
.ca
.map(|index| self.protein.confidence(self.protein.atoms[index].b_factor));
let second = if let Some(confidence) = confidence {
format!(
"pLDDT {:.2} ({}){}",
confidence,
confidence_label(confidence),
if selected {
" click again or x to clear"
} else {
" click for atom detail"
}
)
} else if selected {
"click again or x to clear".to_owned()
} else {
"click for atom detail".to_owned()
};
push_terminal_text(
frame,
card.column + 1,
card.row.saturating_add(1),
truncate(&second, usize::from(card.width_columns.saturating_sub(2))),
ACCENT,
card_background,
false,
);
}
if let Some(atlas) = &self.atlas
&& molecule_columns < columns
{
self.add_terminal_pae_text(frame, atlas, layout, surface);
}
let status = format!(
"{} atoms | {} residues | {} {} | {} | {}{}",
self.protein.atoms.len(),
self.protein.residue_count,
self.protein.chains.len(),
if self.protein.chains.len() == 1 {
"chain"
} else {
"chains"
},
self.representation.name(),
self.color_scheme.compact_name(),
if self.auto_rotate { " | AUTO" } else { "" }
);
push_terminal_text(
frame,
1,
rows.saturating_sub(2),
truncate(&status, usize::from(columns.saturating_sub(2))),
ACCENT,
title_background,
true,
);
let help = "hover inspect click detail drag rotate wheel zoom m mode c color x clear r reset q quit";
push_terminal_text(
frame,
1,
rows.saturating_sub(1),
truncate(help, usize::from(columns.saturating_sub(2))),
MUTED,
title_background,
false,
);
}
fn add_terminal_sequence(
&self,
frame: &mut Frame,
columns: u16,
panel_row: u16,
background: Color,
) {
let active = self.selected_residue.or(self.hovered_residue);
let chain = active
.map(|index| self.protein.residues[index].chain.as_str())
.or_else(|| self.protein.chains.first().map(String::as_str))
.unwrap_or("_");
let residue_indices = self
.protein
.residues
.iter()
.enumerate()
.filter_map(|(index, residue)| {
(residue.chain == chain && residue.ca.is_some()).then_some(index)
})
.collect::<Vec<_>>();
if residue_indices.is_empty() {
return;
}
let header = format!(
"SEQUENCE CHAIN {chain} {} RESIDUES hover inspect / click select",
residue_indices.len()
);
push_terminal_text(
frame,
1,
panel_row,
truncate(&header, usize::from(columns.saturating_sub(2))),
TEXT,
background,
true,
);
let per_row = usize::from(columns.saturating_sub(4).max(10));
let capacity = per_row * 2;
let active_position =
active.and_then(|active| residue_indices.iter().position(|index| *index == active));
let mut start = active_position
.map(|position| position.saturating_sub(capacity / 2))
.unwrap_or(0)
.min(residue_indices.len().saturating_sub(capacity));
start = start / 10 * 10;
let visible = capacity.min(residue_indices.len().saturating_sub(start));
for sequence_row in 0..2 {
let row_start = sequence_row * per_row;
if row_start >= visible {
break;
}
let row_end = (row_start + per_row).min(visible);
let letters = (row_start..row_end)
.map(|offset| {
amino_acid_letter(&self.protein.residues[residue_indices[start + offset]].name)
})
.collect::<String>();
let number_row = panel_row + 1 + sequence_row as u16 * 2;
push_terminal_text(frame, 2, number_row + 1, letters, TEXT, background, false);
for offset in row_start..row_end {
let sequence_position = start + offset;
if !sequence_position.is_multiple_of(10) {
continue;
}
let residue = &self.protein.residues[residue_indices[sequence_position]];
push_terminal_text(
frame,
2 + (offset - row_start) as u16,
number_row,
residue.number.to_string(),
[75, 164, 255, 255],
background,
false,
);
}
}
if let Some(active_position) = active_position
&& active_position >= start
&& active_position < start + visible
{
let offset = active_position - start;
let sequence_row = offset / per_row;
let sequence_column = offset % per_row;
let residue_index = residue_indices[active_position];
let residue_background = if self.selected_residue == Some(residue_index) {
[190, 35, 111, 255]
} else {
[112, 42, 90, 255]
};
push_terminal_text(
frame,
2 + sequence_column as u16,
panel_row + 2 + sequence_row as u16 * 2,
amino_acid_letter(&self.protein.residues[residue_index].name).to_string(),
TEXT,
residue_background,
self.selected_residue == Some(residue_index),
);
}
}
fn add_terminal_pae_text(
&self,
frame: &mut Frame,
atlas: &AtlasPanel,
layout: ViewLayout,
surface: TerminalSurface,
) {
let TerminalSurface {
width,
height,
columns,
rows,
} = surface;
let panel_column = pixel_to_column(layout.molecule_width, width, columns);
let panel_row = pixel_to_row(layout.scene_top, height, rows);
let panel_columns = columns.saturating_sub(panel_column);
let panel_background = PANEL_BACKGROUND;
push_terminal_text(
frame,
panel_column.saturating_add(1),
panel_row,
truncate(
"PREDICTED ALIGNED ERROR",
usize::from(panel_columns.saturating_sub(2)),
),
TEXT,
panel_background,
true,
);
let Some(confidence) = &atlas.confidence else {
push_terminal_text(
frame,
panel_column.saturating_add(1),
panel_row.saturating_add(2),
truncate(&atlas.label, usize::from(panel_columns.saturating_sub(2))),
ACCENT,
panel_background,
true,
);
push_terminal_text(
frame,
panel_column.saturating_add(1),
panel_row.saturating_add(3),
"PAE unavailable for live ESMFold",
MUTED,
panel_background,
false,
);
return;
};
let panel = Rect {
x: layout.molecule_width,
y: layout.scene_top,
width: width.saturating_sub(layout.molecule_width),
height: layout.scene_height,
};
let cell_height = (height / u32::from(rows.max(1))).max(1);
let available_height = panel.height.saturating_sub(cell_height.saturating_mul(8));
let heat_size = panel.width.saturating_sub(28).min(available_height).max(16);
let heat_x = panel.x + (panel.width.saturating_sub(heat_size)) / 2;
let heat_y = panel.y + cell_height * 2;
let last = confidence.pae.len().saturating_sub(1).to_string();
push_terminal_text(
frame,
pixel_to_column(heat_x, width, columns),
pixel_to_row(heat_y.saturating_add(heat_size), height, rows),
"0",
MUTED,
panel_background,
false,
);
let end_column = pixel_to_column(heat_x.saturating_add(heat_size), width, columns)
.saturating_sub(last.len() as u16);
push_terminal_text(
frame,
end_column,
pixel_to_row(heat_y.saturating_add(heat_size), height, rows),
last,
MUTED,
panel_background,
false,
);
let legend_width = heat_size.min(220);
let legend_x = panel.x + (panel.width.saturating_sub(legend_width)) / 2;
let legend_y = heat_y
.saturating_add(heat_size)
.saturating_add(cell_height.saturating_mul(2));
let legend_row = pixel_to_row(
legend_y.saturating_add((cell_height / 3).max(1)),
height,
rows,
);
let legend_start = pixel_to_column(legend_x, width, columns);
let legend_middle =
pixel_to_column(legend_x.saturating_add(legend_width / 2), width, columns)
.saturating_sub(1);
let legend_end = pixel_to_column(legend_x.saturating_add(legend_width), width, columns)
.saturating_sub(2);
for (column, label) in [
(legend_start, "0"),
(legend_middle, "15"),
(legend_end, "30"),
] {
push_terminal_text(
frame,
column,
legend_row,
label,
TEXT,
panel_background,
false,
);
}
let metric = format!("{} pTM {:.3}", atlas.label, confidence.ptm);
push_terminal_text(
frame,
panel_column.saturating_add(1),
pixel_to_row(
panel
.y
.saturating_add(panel.height)
.saturating_sub(cell_height),
height,
rows,
),
truncate(&metric, usize::from(panel_columns.saturating_sub(2))),
ACCENT,
panel_background,
true,
);
}
fn render_selection_detail(
&self,
canvas: &mut Canvas,
residue_index: usize,
projected: &[Projected],
_scene_center: [f32; 2],
scene_scale: f32,
) {
let selected_atoms = self.residue_atom_indices(residue_index);
if selected_atoms.is_empty() {
return;
}
let mut detail_atoms = Vec::new();
for (index, atom) in self.protein.atoms.iter().enumerate() {
if atom.element == "H" {
continue;
}
let is_selected = selected_atoms.contains(&index);
let nearby = !is_selected
&& selected_atoms.iter().any(|selected| {
atom.position
.distance(self.protein.atoms[*selected].position)
<= 4.2
});
if is_selected || nearby {
detail_atoms.push(index);
if detail_atoms.len() >= 160 {
break;
}
}
}
for left in 0..detail_atoms.len() {
for right in left + 1..detail_atoms.len() {
let first_index = detail_atoms[left];
let second_index = detail_atoms[right];
let first = &self.protein.atoms[first_index];
let second = &self.protein.atoms[second_index];
let distance = first.position.distance(second.position);
if distance < 0.45
|| distance
> covalent_radius(&first.element) + covalent_radius(&second.element) + 0.45
{
continue;
}
let from = projected[first_index];
let to = projected[second_index];
let touches_selection =
selected_atoms.contains(&first_index) || selected_atoms.contains(&second_index);
if touches_selection {
canvas.line_3d(
from.array(),
to.array(),
5.5,
[247, 66, 151, 255],
[247, 66, 151, 255],
);
}
canvas.line_3d(
from.array(),
to.array(),
2.8,
element_color(&first.element),
element_color(&second.element),
);
}
}
let mut contacts = 0;
for selected_index in &selected_atoms {
let selected_atom = &self.protein.atoms[*selected_index];
if !is_contact_element(&selected_atom.element) {
continue;
}
for environment_index in &detail_atoms {
if selected_atoms.contains(environment_index) {
continue;
}
let environment = &self.protein.atoms[*environment_index];
if !is_contact_element(&environment.element) {
continue;
}
let distance = selected_atom.position.distance(environment.position);
if !(2.2..=3.6).contains(&distance) {
continue;
}
self.render_dashed_contact(
canvas,
projected[*selected_index],
projected[*environment_index],
);
contacts += 1;
if contacts >= 20 {
break;
}
}
if contacts >= 20 {
break;
}
}
for atom_index in detail_atoms {
let atom = &self.protein.atoms[atom_index];
let point = projected[atom_index];
let radius = (element_radius(&atom.element) / self.protein.radius * scene_scale * 0.42)
.clamp(2.4, 8.0);
if selected_atoms.contains(&atom_index) {
canvas.sphere(
point.x,
point.y,
point.z + 0.003,
radius + 2.0,
[247, 66, 151, 255],
);
}
canvas.sphere(
point.x,
point.y,
point.z + 0.005,
radius,
element_color(&atom.element),
);
}
}
fn render_dashed_contact(&self, canvas: &mut Canvas, from: Projected, to: Projected) {
const DASHES: usize = 7;
for dash in 0..DASHES {
let start = dash as f32 / DASHES as f32;
let end = (dash as f32 + 0.55) / DASHES as f32;
canvas.line_3d(
[
from.x + (to.x - from.x) * start,
from.y + (to.y - from.y) * start,
from.z + (to.z - from.z) * start + 0.004,
],
[
from.x + (to.x - from.x) * end,
from.y + (to.y - from.y) * end,
from.z + (to.z - from.z) * end + 0.004,
],
1.8,
[48, 170, 225, 255],
[48, 170, 225, 255],
);
}
}
fn residue_atom_indices(&self, residue_index: usize) -> Vec<usize> {
let residue = &self.protein.residues[residue_index];
self.protein
.atoms
.iter()
.enumerate()
.filter_map(|(index, atom)| {
(atom.chain == residue.chain
&& atom.residue_number == residue.number
&& atom.insertion_code == residue.insertion_code)
.then_some(index)
})
.collect()
}
fn pick_residue(
&self,
x: u32,
y: u32,
width: u32,
height: u32,
terminal: Option<(u16, u16)>,
) -> Option<usize> {
let layout = terminal.map_or_else(
|| self.view_layout(width, height),
|(columns, rows)| self.terminal_view_layout(width, height, columns, rows),
);
if x >= layout.molecule_width
|| y < layout.scene_top
|| y >= layout.scene_top.saturating_add(layout.scene_height)
{
return None;
}
let threshold = (1.25 / self.protein.radius * layout.scene_scale + 9.0).clamp(11.0, 30.0);
let mut best: Option<(usize, f32)> = None;
for (index, residue) in self.protein.residues.iter().enumerate() {
let Some(ca) = residue.ca else {
continue;
};
let point = self.project_position(
self.protein.atoms[ca].position,
layout.scene_center,
layout.scene_scale,
);
let dx = point.x - x as f32;
let dy = point.y - y as f32;
let distance = (dx * dx + dy * dy).sqrt();
if distance > threshold {
continue;
}
let score = distance - point.z * 3.0;
if best.is_none_or(|(_, best_score)| score < best_score) {
best = Some((index, score));
}
}
best.map(|(index, _)| index)
}
fn set_selection(&mut self, residue: Option<usize>) {
if residue.is_some() && residue == self.selected_residue {
self.selected_residue = None;
self.view_center = self.protein.center;
self.zoom = 1.35;
return;
}
self.selected_residue = residue;
if let Some(index) = residue
&& let Some(ca) = self.protein.residues[index].ca
{
self.view_center = self.protein.atoms[ca].position;
self.zoom = (self.protein.radius / 9.0).clamp(2.6, 6.0);
self.auto_rotate = false;
} else {
self.view_center = self.protein.center;
self.zoom = 1.35;
}
}
fn frame_dimensions(context: &AppContext) -> (u32, u32) {
let scale = if context.quality == "fast"
&& matches!(context.protocol.as_str(), "kitty" | "iterm2")
{
0.58
} else {
1.0
};
(
((context.geometry.pixel_width as f32 * scale).round() as u32).max(1),
((context.geometry.pixel_height as f32 * scale).round() as u32).max(1),
)
}
fn scale_mouse_point(
x: u32,
y: u32,
context: &AppContext,
width: u32,
height: u32,
) -> (u32, u32) {
(
x.saturating_mul(width) / context.geometry.pixel_width.max(1),
y.saturating_mul(height) / context.geometry.pixel_height.max(1),
)
}
fn render_ribbon(
&self,
canvas: &mut Canvas,
scene_center: [f32; 2],
scene_scale: f32,
fast: bool,
) {
let subdivisions = if fast { 4 } else { 8 };
let radial_segments = if fast { 8 } else { 12 };
for chain in ribbon_chains(&self.protein) {
let samples = build_ribbon_samples(&self.protein, &chain, subdivisions);
if samples.is_empty() {
continue;
}
let mut start = 0;
while start < samples.len() {
let secondary = samples[start].secondary;
let mut end = start + 1;
while end < samples.len() && samples[end].secondary == secondary {
end += 1;
}
let draw_end = if end < samples.len() { end + 1 } else { end };
let run = &samples[start..draw_end];
if run.len() == 1 {
let point = self.project_position(run[0].center, scene_center, scene_scale);
let radius =
(run[0].half_width / self.protein.radius * scene_scale).clamp(1.0, 8.0);
canvas.sphere(point.x, point.y, point.z, radius, self.sample_color(run[0]));
} else if secondary == SecondaryStructure::Sheet {
self.render_box_run(
canvas,
run,
scene_center,
scene_scale,
start == 0,
end == samples.len(),
);
} else {
self.render_elliptical_run(
canvas,
run,
radial_segments,
scene_center,
scene_scale,
[start == 0, end == samples.len()],
);
}
start = end;
}
}
}
fn render_elliptical_run(
&self,
canvas: &mut Canvas,
samples: &[RibbonSample],
radial_segments: usize,
scene_center: [f32; 2],
scene_scale: f32,
caps: [bool; 2],
) {
for pair in samples.windows(2) {
for segment in 0..radial_segments {
let next = (segment + 1) % radial_segments;
let (a, normal_a) = ellipse_vertex(pair[0], segment, radial_segments);
let (b, normal_b) = ellipse_vertex(pair[0], next, radial_segments);
let (c, normal_c) = ellipse_vertex(pair[1], segment, radial_segments);
let (d, normal_d) = ellipse_vertex(pair[1], next, radial_segments);
self.draw_mesh_triangle(
canvas,
[
(a, normal_a, pair[0]),
(c, normal_c, pair[1]),
(b, normal_b, pair[0]),
],
scene_center,
scene_scale,
);
self.draw_mesh_triangle(
canvas,
[
(b, normal_b, pair[0]),
(c, normal_c, pair[1]),
(d, normal_d, pair[1]),
],
scene_center,
scene_scale,
);
}
}
if caps[0] {
self.render_ellipse_cap(
canvas,
samples[0],
radial_segments,
-samples[0].tangent,
scene_center,
scene_scale,
);
}
if caps[1] {
let sample = *samples.last().expect("non-empty ribbon run");
self.render_ellipse_cap(
canvas,
sample,
radial_segments,
sample.tangent,
scene_center,
scene_scale,
);
}
}
fn render_ellipse_cap(
&self,
canvas: &mut Canvas,
sample: RibbonSample,
radial_segments: usize,
normal: Vec3,
scene_center: [f32; 2],
scene_scale: f32,
) {
for segment in 0..radial_segments {
let next = (segment + 1) % radial_segments;
let (first, _) = ellipse_vertex(sample, segment, radial_segments);
let (second, _) = ellipse_vertex(sample, next, radial_segments);
self.draw_mesh_triangle(
canvas,
[
(sample.center, normal, sample),
(first, normal, sample),
(second, normal, sample),
],
scene_center,
scene_scale,
);
}
}
fn render_box_run(
&self,
canvas: &mut Canvas,
samples: &[RibbonSample],
scene_center: [f32; 2],
scene_scale: f32,
start_cap: bool,
end_cap: bool,
) {
for pair in samples.windows(2) {
for face in 0..4 {
let next = (face + 1) % 4;
let a = box_vertex(pair[0], face);
let b = box_vertex(pair[0], next);
let c = box_vertex(pair[1], face);
let d = box_vertex(pair[1], next);
let normal = (c - a).cross(b - a).normalized();
self.draw_mesh_triangle(
canvas,
[
(a, normal, pair[0]),
(c, normal, pair[1]),
(b, normal, pair[0]),
],
scene_center,
scene_scale,
);
self.draw_mesh_triangle(
canvas,
[
(b, normal, pair[0]),
(c, normal, pair[1]),
(d, normal, pair[1]),
],
scene_center,
scene_scale,
);
}
}
if start_cap {
self.render_box_cap(
canvas,
samples[0],
-samples[0].tangent,
scene_center,
scene_scale,
);
}
if end_cap {
let sample = *samples.last().expect("non-empty ribbon run");
self.render_box_cap(canvas, sample, sample.tangent, scene_center, scene_scale);
}
}
fn render_box_cap(
&self,
canvas: &mut Canvas,
sample: RibbonSample,
normal: Vec3,
scene_center: [f32; 2],
scene_scale: f32,
) {
for face in 0..4 {
let first = box_vertex(sample, face);
let second = box_vertex(sample, (face + 1) % 4);
self.draw_mesh_triangle(
canvas,
[
(sample.center, normal, sample),
(first, normal, sample),
(second, normal, sample),
],
scene_center,
scene_scale,
);
}
}
fn draw_mesh_triangle(
&self,
canvas: &mut Canvas,
vertices: [(Vec3, Vec3, RibbonSample); 3],
scene_center: [f32; 2],
scene_scale: f32,
) {
let projected = vertices
.map(|(position, _, _)| self.project_position(position, scene_center, scene_scale));
let colors = vertices.map(|(_, normal, sample)| {
scale(self.sample_color(sample), self.surface_light(normal))
});
canvas.triangle_3d_gradient(
projected[0].array(),
projected[1].array(),
projected[2].array(),
colors[0],
colors[1],
colors[2],
);
}
fn surface_light(&self, normal: Vec3) -> f32 {
let view_normal = rotate(normal, self.yaw, self.pitch).normalized();
let light = Vec3::new(-0.35, 0.45, 0.82).normalized();
let diffuse = view_normal.dot(light).max(0.0);
let soft_fill = view_normal.z.abs();
let half_vector = (light + Vec3::new(0.0, 0.0, 1.0)).normalized();
let specular = view_normal.dot(half_vector).max(0.0).powi(18);
(0.56 + diffuse * 0.34 + soft_fill * 0.14 + specular * 0.12).clamp(0.5, 1.12)
}
fn project_position(
&self,
position: Vec3,
scene_center: [f32; 2],
scene_scale: f32,
) -> Projected {
let local = (position - self.view_center) * (1.0 / self.protein.radius);
let rotated = rotate(local, self.yaw, self.pitch);
Projected {
x: scene_center[0] + rotated.x * scene_scale,
y: scene_center[1] - rotated.y * scene_scale,
z: rotated.z,
}
}
fn sample_color(&self, sample: RibbonSample) -> Color {
if self.selected_residue == Some(sample.residue_index) {
return [247, 66, 151, 255];
}
if self.hovered_residue == Some(sample.residue_index) {
return [255, 112, 174, 255];
}
let atom = &self.protein.atoms[sample.atom_index];
let local = (sample.center - self.view_center) * (1.0 / self.protein.radius);
let z = rotate(local, self.yaw, self.pitch).z;
self.atom_color(atom, z)
}
fn atom_color(&self, atom: &Atom, z: f32) -> Color {
let base = match self.color_scheme {
ColorScheme::Chain => chain_color(&atom.chain, &self.protein.chains),
ColorScheme::Element => element_color(&atom.element),
ColorScheme::Confidence => confidence_color(self.protein.confidence(atom.b_factor)),
};
scale(base, 0.82 + (z * 0.5 + 0.5).clamp(0.0, 1.0) * 0.18)
}
}
impl App for Viewer {
fn event(&mut self, event: AppEvent, context: &AppContext) -> Result<ControlFlow, String> {
let render = match event {
AppEvent::Start | AppEvent::Resize { .. } => true,
AppEvent::Key { key, .. } => match key.as_str() {
"left" | "h" => {
self.yaw -= 0.12;
true
}
"right" | "l" => {
self.yaw += 0.12;
true
}
"up" | "k" => {
self.pitch = (self.pitch + 0.12).clamp(-1.55, 1.55);
true
}
"down" | "j" => {
self.pitch = (self.pitch - 0.12).clamp(-1.55, 1.55);
true
}
"+" | "=" => {
self.zoom = (self.zoom * 1.12).min(8.0);
true
}
"-" | "_" => {
self.zoom = (self.zoom / 1.12).max(0.15);
true
}
"m" => {
self.representation = self.representation.next();
true
}
"c" => {
self.color_scheme = self.color_scheme.next();
true
}
" " | "space" => {
self.auto_rotate = !self.auto_rotate;
true
}
"r" => {
self.reset_view();
true
}
"x" => {
self.set_selection(None);
true
}
_ => false,
},
AppEvent::Mouse {
kind,
column,
row,
pixel_x,
pixel_y,
button,
} => {
let raw_x = pixel_x.unwrap_or(u32::from(column) * context.geometry.cells.width);
let raw_y = pixel_y.unwrap_or(u32::from(row) * context.geometry.cells.height);
let (width, height) = Self::frame_dimensions(context);
let (x, y) = Self::scale_mouse_point(raw_x, raw_y, context, width, height);
match (kind.as_str(), button.as_deref()) {
("move", _) => {
let hovered = self.pick_residue(
x,
y,
width,
height,
Some((
context.geometry.content.width,
context.geometry.content.height,
)),
);
if hovered != self.hovered_residue {
self.hovered_residue = hovered;
true
} else {
false
}
}
("down", Some("left")) => {
self.dragging = Some((x, y));
self.dragged = false;
self.click_candidate = self.pick_residue(
x,
y,
width,
height,
Some((
context.geometry.content.width,
context.geometry.content.height,
)),
);
let changed = self.hovered_residue != self.click_candidate;
self.hovered_residue = self.click_candidate;
changed
}
("drag", Some("left")) => {
self.dragged = true;
self.click_candidate = None;
if let Some((last_x, last_y)) = self.dragging.replace((x, y)) {
self.yaw += (x as f32 - last_x as f32) * 0.008;
self.pitch = (self.pitch - (y as f32 - last_y as f32) * 0.008)
.clamp(-1.55, 1.55);
true
} else {
false
}
}
("up", Some("left")) => {
self.dragging = None;
let picked = if self.dragged {
None
} else {
self.pick_residue(
x,
y,
width,
height,
Some((
context.geometry.content.width,
context.geometry.content.height,
)),
)
};
self.click_candidate = None;
let was_dragged = self.dragged;
self.dragged = false;
if was_dragged {
false
} else {
self.set_selection(picked);
true
}
}
("scroll_up", _) => {
self.zoom = (self.zoom * 1.12).min(8.0);
true
}
("scroll_down", _) => {
self.zoom = (self.zoom / 1.12).max(0.15);
true
}
_ => false,
}
}
AppEvent::Tick { elapsed_ms } if self.auto_rotate => {
let delta = elapsed_ms.saturating_sub(self.last_tick_ms).min(100);
self.last_tick_ms = elapsed_ms;
self.yaw += delta as f32 * 0.00055;
true
}
AppEvent::Tick { elapsed_ms } => {
self.last_tick_ms = elapsed_ms;
false
}
};
Ok(if render {
ControlFlow::Render
} else {
ControlFlow::Continue
})
}
fn render(&mut self, context: &AppContext) -> Result<Frame, String> {
let (width, height) = Self::frame_dimensions(context);
Ok(self.render_scene(
width,
height,
context.quality == "fast",
Some((
context.geometry.content.width,
context.geometry.content.height,
)),
))
}
}
#[derive(Clone, Copy)]
struct Projected {
x: f32,
y: f32,
z: f32,
}
impl Projected {
const fn array(self) -> [f32; 3] {
[self.x, self.y, self.z]
}
}
#[derive(Clone, Copy)]
struct RibbonSample {
center: Vec3,
tangent: Vec3,
normal: Vec3,
binormal: Vec3,
half_width: f32,
half_thickness: f32,
secondary: SecondaryStructure,
atom_index: usize,
residue_index: usize,
}
#[derive(Clone, Copy)]
struct RibbonNode {
center: Vec3,
normal: Vec3,
half_width: f32,
half_thickness: f32,
secondary: SecondaryStructure,
atom_index: usize,
residue_index: usize,
}
#[derive(Clone, Copy)]
struct Rect {
x: u32,
y: u32,
width: u32,
height: u32,
}
#[derive(Clone, Copy)]
struct ViewLayout {
detailed: bool,
font_scale: u32,
title_height: u32,
footer_height: u32,
sequence_height: u32,
scene_top: u32,
scene_height: u32,
molecule_width: u32,
scene_center: [f32; 2],
scene_scale: f32,
}
#[derive(Clone, Copy)]
struct TerminalSurface {
width: u32,
height: u32,
columns: u16,
rows: u16,
}
#[derive(Clone, Copy)]
struct TerminalCardLayout {
column: u16,
row: u16,
width_columns: u16,
rect: Rect,
}
fn terminal_card_layout(layout: ViewLayout, surface: TerminalSurface) -> TerminalCardLayout {
let molecule_columns =
pixel_to_column_edge(layout.molecule_width, surface.width, surface.columns).max(1);
let width_columns = molecule_columns.saturating_sub(2).clamp(20, 64);
let column = molecule_columns.saturating_sub(width_columns + 1);
let scene_bottom = pixel_to_row_edge(
layout.scene_top.saturating_add(layout.scene_height),
surface.height,
surface.rows,
);
let row = scene_bottom.saturating_sub(3);
let x = column_to_pixel_edge(column, surface.width, surface.columns);
let right = column_to_pixel_edge(
column.saturating_add(width_columns),
surface.width,
surface.columns,
);
let y = row_to_pixel_edge(row, surface.height, surface.rows);
let bottom = row_to_pixel_edge(scene_bottom, surface.height, surface.rows);
TerminalCardLayout {
column,
row,
width_columns,
rect: Rect {
x,
y,
width: right.saturating_sub(x).max(1),
height: bottom.saturating_sub(y).max(1),
},
}
}
fn pixel_to_column(x: u32, width: u32, columns: u16) -> u16 {
((u64::from(x) * u64::from(columns) / u64::from(width.max(1)))
.min(u64::from(columns.saturating_sub(1)))) as u16
}
fn pixel_to_column_edge(x: u32, width: u32, columns: u16) -> u16 {
(u64::from(x) * u64::from(columns) / u64::from(width.max(1))).min(u64::from(columns)) as u16
}
fn pixel_to_row(y: u32, height: u32, rows: u16) -> u16 {
((u64::from(y) * u64::from(rows) / u64::from(height.max(1)))
.min(u64::from(rows.saturating_sub(1)))) as u16
}
fn pixel_to_row_edge(y: u32, height: u32, rows: u16) -> u16 {
(u64::from(y) * u64::from(rows) / u64::from(height.max(1))).min(u64::from(rows)) as u16
}
fn column_to_pixel_edge(column: u16, width: u32, columns: u16) -> u32 {
(u64::from(column) * u64::from(width) / u64::from(columns.max(1))) as u32
}
fn row_to_pixel_edge(row: u16, height: u32, rows: u16) -> u32 {
(u64::from(row) * u64::from(height) / u64::from(rows.max(1))) as u32
}
fn push_terminal_text(
frame: &mut Frame,
column: u16,
row: u16,
text: impl Into<String>,
foreground: Color,
background: Color,
bold: bool,
) {
let mut run = TextRun::new(column, row, text)
.foreground(rgba(foreground))
.background(rgba(background));
if bold {
run = run.bold();
}
frame.push_text(run);
}
const fn rgba(color: Color) -> Rgba {
Rgba {
r: color[0],
g: color[1],
b: color[2],
a: color[3],
}
}
fn ribbon_chains(protein: &Protein) -> Vec<Vec<usize>> {
let mut chains = Vec::new();
let mut current = Vec::new();
for (index, residue) in protein.residues.iter().enumerate() {
let Some(ca) = residue.ca else {
if !current.is_empty() {
chains.push(std::mem::take(&mut current));
}
continue;
};
let connected = current.last().is_some_and(|previous: &usize| {
let previous_residue = &protein.residues[*previous];
previous_residue.chain == residue.chain
&& previous_residue.ca.is_some_and(|previous_ca| {
protein.atoms[previous_ca]
.position
.distance(protein.atoms[ca].position)
<= 4.5
})
});
if !connected && !current.is_empty() {
chains.push(std::mem::take(&mut current));
}
current.push(index);
}
if !current.is_empty() {
chains.push(current);
}
chains
}
fn build_ribbon_samples(
protein: &Protein,
residue_indices: &[usize],
subdivisions: usize,
) -> Vec<RibbonSample> {
let mut nodes = Vec::with_capacity(residue_indices.len());
let mut previous_normal: Option<Vec3> = None;
for (position, residue_index) in residue_indices.iter().copied().enumerate() {
let residue = &protein.residues[residue_index];
let atom_index = residue
.ca
.expect("ribbon chain only contains C-alpha atoms");
let center = protein.atoms[atom_index].position;
let previous = position
.checked_sub(1)
.and_then(|index| protein.residues[residue_indices[index]].ca)
.map(|index| protein.atoms[index].position)
.unwrap_or(center);
let next = residue_indices
.get(position + 1)
.and_then(|index| protein.residues[*index].ca)
.map(|index| protein.atoms[index].position)
.unwrap_or(center);
let tangent = (next - previous).normalized();
let guide = residue
.c
.zip(residue.o)
.map(|(carbon, oxygen)| protein.atoms[oxygen].position - protein.atoms[carbon].position)
.or_else(|| {
residue
.o
.map(|index| protein.atoms[index].position - center)
})
.unwrap_or_else(|| tangent.cross(Vec3::new(0.0, 0.0, 1.0)));
let mut normal = orthogonalized(guide, tangent);
if normal.length() < 0.1 {
normal = orthogonalized(Vec3::new(0.0, 1.0, 0.0), tangent);
}
if previous_normal.is_some_and(|previous| previous.dot(normal) < 0.0) {
normal = -normal;
}
previous_normal = Some(normal);
let (half_width, half_thickness) = match residue.secondary {
SecondaryStructure::Coil => (0.30, 0.30),
SecondaryStructure::Helix => (1.02, 0.22),
SecondaryStructure::Sheet => (1.16, 0.14),
};
nodes.push(RibbonNode {
center,
normal,
half_width,
half_thickness,
secondary: residue.secondary,
atom_index,
residue_index,
});
}
add_sheet_arrowheads(&mut nodes);
if nodes.len() == 1 {
let tangent = Vec3::new(0.0, 0.0, 1.0);
return vec![RibbonSample {
center: nodes[0].center,
tangent,
normal: nodes[0].normal,
binormal: tangent.cross(nodes[0].normal).normalized(),
half_width: nodes[0].half_width,
half_thickness: nodes[0].half_thickness,
secondary: nodes[0].secondary,
atom_index: nodes[0].atom_index,
residue_index: nodes[0].residue_index,
}];
}
let subdivisions = subdivisions.max(1);
let mut samples = Vec::with_capacity((nodes.len() - 1) * subdivisions + 1);
for index in 0..nodes.len() - 1 {
let first = nodes[index.saturating_sub(1)];
let from = nodes[index];
let to = nodes[index + 1];
let fourth = nodes[(index + 2).min(nodes.len() - 1)];
for step in 0..subdivisions {
let amount = step as f32 / subdivisions as f32;
samples.push(RibbonSample {
center: catmull_rom(first.center, from.center, to.center, fourth.center, amount),
tangent: Vec3::default(),
normal: from.normal.lerp(to.normal, smoothstep(amount)).normalized(),
binormal: Vec3::default(),
half_width: from.half_width + (to.half_width - from.half_width) * amount,
half_thickness: from.half_thickness
+ (to.half_thickness - from.half_thickness) * amount,
secondary: if amount < 0.5 {
from.secondary
} else {
to.secondary
},
atom_index: if amount < 0.5 {
from.atom_index
} else {
to.atom_index
},
residue_index: if amount < 0.5 {
from.residue_index
} else {
to.residue_index
},
});
}
}
let last = *nodes.last().expect("non-empty ribbon");
samples.push(RibbonSample {
center: last.center,
tangent: Vec3::default(),
normal: last.normal,
binormal: Vec3::default(),
half_width: last.half_width,
half_thickness: last.half_thickness,
secondary: last.secondary,
atom_index: last.atom_index,
residue_index: last.residue_index,
});
finish_ribbon_frames(&mut samples);
samples
}
fn add_sheet_arrowheads(nodes: &mut [RibbonNode]) {
let mut start = 0;
while start < nodes.len() {
if nodes[start].secondary != SecondaryStructure::Sheet {
start += 1;
continue;
}
let mut end = start + 1;
while end < nodes.len() && nodes[end].secondary == SecondaryStructure::Sheet {
end += 1;
}
if end - start >= 3 {
nodes[end - 2].half_width *= 1.55;
nodes[end - 1].half_width = 0.08;
}
start = end;
}
}
fn finish_ribbon_frames(samples: &mut [RibbonSample]) {
let mut previous_normal: Option<Vec3> = None;
for index in 0..samples.len() {
let previous = index
.checked_sub(1)
.map(|previous| samples[previous].center)
.unwrap_or(samples[index].center);
let next = samples
.get(index + 1)
.map(|next| next.center)
.unwrap_or(samples[index].center);
let tangent = (next - previous).normalized();
let mut normal = orthogonalized(samples[index].normal, tangent);
if normal.length() < 0.1 {
normal = orthogonalized(Vec3::new(0.0, 1.0, 0.0), tangent);
}
if previous_normal.is_some_and(|previous| previous.dot(normal) < 0.0) {
normal = -normal;
}
let binormal = tangent.cross(normal).normalized();
samples[index].tangent = tangent;
samples[index].normal = normal;
samples[index].binormal = binormal;
previous_normal = Some(normal);
}
if samples.len() > 4 {
let normals: Vec<Vec3> = samples.iter().map(|sample| sample.normal).collect();
for index in 1..samples.len() - 1 {
let blended =
(normals[index - 1] + normals[index] * 2.0 + normals[index + 1]).normalized();
let normal = orthogonalized(blended, samples[index].tangent);
samples[index].normal = normal;
samples[index].binormal = samples[index].tangent.cross(normal).normalized();
}
}
}
fn orthogonalized(vector: Vec3, tangent: Vec3) -> Vec3 {
(vector - tangent * vector.dot(tangent)).normalized()
}
fn smoothstep(amount: f32) -> f32 {
amount * amount * (3.0 - 2.0 * amount)
}
fn ellipse_vertex(sample: RibbonSample, segment: usize, radial_segments: usize) -> (Vec3, Vec3) {
let angle = segment as f32 / radial_segments as f32 * std::f32::consts::TAU;
let (sin, cos) = angle.sin_cos();
let position = sample.center
+ sample.normal * (sample.half_width * cos)
+ sample.binormal * (sample.half_thickness * sin);
let normal = (sample.normal * (sample.half_thickness * cos)
+ sample.binormal * (sample.half_width * sin))
.normalized();
(position, normal)
}
fn box_vertex(sample: RibbonSample, corner: usize) -> Vec3 {
let (width, thickness) = match corner % 4 {
0 => (sample.half_width, sample.half_thickness),
1 => (-sample.half_width, sample.half_thickness),
2 => (-sample.half_width, -sample.half_thickness),
_ => (sample.half_width, -sample.half_thickness),
};
sample.center + sample.normal * width + sample.binormal * thickness
}
fn catmull_rom(first: Vec3, from: Vec3, to: Vec3, fourth: Vec3, amount: f32) -> Vec3 {
let squared = amount * amount;
let cubed = squared * amount;
(from * 2.0
+ (to - first) * amount
+ (first * 2.0 - from * 5.0 + to * 4.0 - fourth) * squared
+ (-first + from * 3.0 - to * 3.0 + fourth) * cubed)
* 0.5
}
fn render_pae_panel(
canvas: &mut Canvas,
panel: Rect,
atlas: &AtlasPanel,
draw_text: bool,
cell_width: u32,
cell_height: u32,
) {
canvas.fill_rect(
panel.x,
panel.y,
panel.width,
panel.height,
PANEL_BACKGROUND,
);
canvas.fill_rect(panel.x, panel.y, 1, panel.height, [49, 60, 78, 255]);
if panel.width < 80 || panel.height < 80 {
return;
}
let title = "PREDICTED ALIGNED ERROR";
if draw_text {
canvas.text(
panel.x + 10,
panel.y + 8,
&truncate(
title,
panel.width.saturating_sub(20) as usize / cell_width as usize,
),
TEXT,
);
}
let Some(confidence) = &atlas.confidence else {
if draw_text {
let message_y = panel.y + cell_height + 20;
canvas.text(panel.x + 10, message_y, &atlas.label, ACCENT);
canvas.text(
panel.x + 10,
message_y + cell_height + 8,
"PAE is unavailable for",
MUTED,
);
canvas.text(
panel.x + 10,
message_y + (cell_height + 8) * 2,
"live ESMFold predictions.",
MUTED,
);
}
return;
};
let bottom_space = if draw_text {
cell_height.saturating_mul(3) + 72
} else {
cell_height.saturating_mul(8)
};
let available_height = panel.height.saturating_sub(bottom_space);
let heat_size = panel.width.saturating_sub(28).min(available_height).max(16);
let heat_x = panel.x + (panel.width.saturating_sub(heat_size)) / 2;
let heat_y = panel.y
+ if draw_text {
cell_height + 18
} else {
cell_height * 2
};
let rows = confidence.pae.len();
let columns = confidence.pae.first().map_or(0, Vec::len);
for y in 0..heat_size {
let row = (y as usize * rows / heat_size as usize).min(rows.saturating_sub(1));
for x in 0..heat_size {
let column = (x as usize * columns / heat_size as usize).min(columns.saturating_sub(1));
canvas.set_pixel(
heat_x + x,
heat_y + y,
pae_color(confidence.pae[row][column]),
);
}
}
canvas.stroke_rect(heat_x, heat_y, heat_size, heat_size, [126, 143, 160, 255]);
let last = rows.saturating_sub(1);
if draw_text {
canvas.text(heat_x, heat_y + heat_size + 4, "0", MUTED);
}
let end_label = last.to_string();
if draw_text {
canvas.text(
heat_x.saturating_add(heat_size.saturating_sub(end_label.len() as u32 * cell_width)),
heat_y + heat_size + 4,
&end_label,
MUTED,
);
}
let legend_y = heat_y + heat_size + cell_height + if draw_text { 10 } else { cell_height };
let legend_width = heat_size.min(220);
let legend_x = panel.x + (panel.width.saturating_sub(legend_width)) / 2;
let legend_height = if draw_text {
6 + cell_height / 4
} else {
(cell_height / 3).max(1)
};
for x in 0..legend_width {
let value = x as f32 / legend_width.saturating_sub(1).max(1) as f32 * 30.0;
canvas.fill_rect(legend_x + x, legend_y, 1, legend_height, pae_color(value));
}
let legend_text_y = legend_y + legend_height + 4;
if draw_text {
canvas.text(legend_x, legend_text_y, "0", TEXT);
canvas.text(
legend_x + legend_width / 2 - cell_width,
legend_text_y,
"15",
TEXT,
);
canvas.text(
legend_x + legend_width.saturating_sub(cell_width * 2),
legend_text_y,
"30",
TEXT,
);
}
let metric = format!("{} pTM {:.3}", atlas.label, confidence.ptm);
if draw_text {
canvas.text(
panel.x + 10,
panel.y + panel.height.saturating_sub(cell_height + 5),
&truncate(
&metric,
panel.width.saturating_sub(20) as usize / cell_width as usize,
),
ACCENT,
);
}
}
fn pae_color(value: f32) -> Color {
if value <= 15.0 {
mix(
[0, 118, 32, 255],
[111, 196, 100, 255],
value.max(0.0) / 15.0,
)
} else {
mix(
[111, 196, 100, 255],
[242, 246, 238, 255],
(value - 15.0) / 15.0,
)
}
}
fn rotate(point: Vec3, yaw: f32, pitch: f32) -> Vec3 {
let (sin_yaw, cos_yaw) = yaw.sin_cos();
let x = point.x * cos_yaw + point.z * sin_yaw;
let z = -point.x * sin_yaw + point.z * cos_yaw;
let (sin_pitch, cos_pitch) = pitch.sin_cos();
Vec3::new(
x,
point.y * cos_pitch - z * sin_pitch,
point.y * sin_pitch + z * cos_pitch,
)
}
fn chain_color(chain: &str, chains: &[String]) -> Color {
const PALETTE: &[Color] = &[
[73, 215, 196, 255],
[103, 155, 255, 255],
[240, 128, 154, 255],
[245, 192, 83, 255],
[177, 132, 245, 255],
[108, 214, 112, 255],
[244, 139, 70, 255],
[87, 200, 238, 255],
];
let index = chains.iter().position(|item| item == chain).unwrap_or(0);
PALETTE[index % PALETTE.len()]
}
fn element_color(element: &str) -> Color {
match element {
"H" => [235, 239, 245, 255],
"C" => [162, 174, 190, 255],
"N" => [76, 127, 255, 255],
"O" => [244, 82, 82, 255],
"S" => [250, 207, 65, 255],
"P" => [255, 146, 54, 255],
"F" | "CL" => [84, 214, 111, 255],
"BR" => [165, 72, 60, 255],
"I" => [148, 79, 201, 255],
"FE" => [224, 102, 51, 255],
"MG" | "CA" | "ZN" | "MN" | "CU" | "CO" | "NI" => [92, 217, 199, 255],
_ => [225, 120, 205, 255],
}
}
fn confidence_color(value: f32) -> Color {
if value >= 90.0 {
[33, 79, 186, 255]
} else if value >= 70.0 {
mix(
[41, 167, 225, 255],
[33, 79, 186, 255],
(value - 70.0) / 20.0,
)
} else if value >= 50.0 {
mix(
[246, 201, 68, 255],
[41, 167, 225, 255],
(value - 50.0) / 20.0,
)
} else {
mix(
[239, 74, 77, 255],
[246, 201, 68, 255],
value.clamp(0.0, 50.0) / 50.0,
)
}
}
fn confidence_label(value: f32) -> &'static str {
if value >= 90.0 {
"very high"
} else if value >= 70.0 {
"confident"
} else if value >= 50.0 {
"low"
} else {
"very low"
}
}
fn amino_acid_letter(name: &str) -> char {
match name {
"ALA" => 'A',
"ARG" => 'R',
"ASN" => 'N',
"ASP" => 'D',
"CYS" => 'C',
"GLN" => 'Q',
"GLU" => 'E',
"GLY" => 'G',
"HIS" => 'H',
"ILE" => 'I',
"LEU" => 'L',
"LYS" => 'K',
"MET" => 'M',
"PHE" => 'F',
"PRO" => 'P',
"SER" => 'S',
"THR" => 'T',
"TRP" => 'W',
"TYR" => 'Y',
"VAL" => 'V',
"SEC" => 'U',
"PYL" => 'O',
"ASX" => 'B',
"GLX" => 'Z',
_ => 'X',
}
}
fn element_radius(element: &str) -> f32 {
match element {
"H" => 1.20,
"C" => 1.70,
"N" => 1.55,
"O" => 1.52,
"F" => 1.47,
"P" => 1.80,
"S" => 1.80,
"CL" => 1.75,
"FE" | "MG" | "CA" | "ZN" | "MN" | "CU" | "CO" | "NI" => 1.65,
_ => 1.70,
}
}
fn covalent_radius(element: &str) -> f32 {
match element {
"H" => 0.31,
"C" => 0.76,
"N" => 0.71,
"O" => 0.66,
"F" => 0.57,
"P" => 1.07,
"S" => 1.05,
"CL" => 1.02,
"BR" => 1.20,
"I" => 1.39,
"FE" => 1.32,
"MG" => 1.41,
"CA" => 1.76,
"ZN" => 1.22,
_ => 0.85,
}
}
fn is_contact_element(element: &str) -> bool {
matches!(element, "N" | "O" | "S")
}
fn truncate(value: &str, maximum: usize) -> String {
if value.chars().count() <= maximum {
return value.to_owned();
}
if maximum <= 3 {
return value.chars().take(maximum).collect();
}
let mut shortened: String = value.chars().take(maximum - 3).collect();
shortened.push_str("...");
shortened
}
#[cfg(test)]
mod tests {
use super::*;
const PICK_PDB: &str = "TITLE ESMFOLD PICK TEST\nATOM 1 N ALA A 1 -1.000 0.000 0.000 1.00 95.00 N \nATOM 2 CA ALA A 1 0.000 0.000 0.000 1.00 95.00 C \nATOM 3 C ALA A 1 1.300 0.000 0.000 1.00 95.00 C \nATOM 4 O ALA A 1 1.900 1.000 0.000 1.00 95.00 O \nATOM 5 CA GLY A 2 3.800 0.000 0.000 1.00 80.00 C \nEND\n";
#[test]
fn rotation_preserves_length() {
let point = Vec3::new(1.0, 2.0, 3.0);
let rotated = rotate(point, 0.7, -0.4);
assert!((point.length() - rotated.length()).abs() < 0.0001);
}
#[test]
fn confidence_palette_has_expected_extremes() {
assert_eq!(confidence_color(95.0), [33, 79, 186, 255]);
assert_eq!(confidence_color(0.0), [239, 74, 77, 255]);
}
#[test]
fn pae_palette_runs_from_green_to_white() {
assert_eq!(pae_color(0.0), [0, 118, 32, 255]);
assert_eq!(pae_color(30.0), [242, 246, 238, 255]);
}
#[test]
fn catmull_rom_passes_through_segment_endpoints() {
let first = Vec3::new(-1.0, 0.0, 0.0);
let from = Vec3::new(0.0, 1.0, 0.0);
let to = Vec3::new(1.0, 1.0, 0.0);
let fourth = Vec3::new(2.0, 0.0, 0.0);
assert_eq!(catmull_rom(first, from, to, fourth, 0.0), from);
assert_eq!(catmull_rom(first, from, to, fourth, 1.0), to);
}
#[test]
fn picks_projected_residue_and_focuses_selection() {
let protein = Protein::from_pdb_str(PICK_PDB, "pick.pdb").unwrap();
let mut viewer = Viewer::new(protein, Representation::Backbone, ColorScheme::Confidence);
let layout = viewer.view_layout(800, 600);
let ca = viewer.protein.residues[0].ca.unwrap();
let projected = viewer.project_position(
viewer.protein.atoms[ca].position,
layout.scene_center,
layout.scene_scale,
);
assert_eq!(
viewer.pick_residue(projected.x as u32, projected.y as u32, 800, 600, None),
Some(0)
);
viewer.set_selection(Some(0));
assert_eq!(viewer.selected_residue, Some(0));
assert_eq!(viewer.view_center, viewer.protein.atoms[ca].position);
assert!(viewer.zoom > 1.35);
}
#[test]
fn maps_standard_amino_acids_to_sequence_letters() {
assert_eq!(amino_acid_letter("HIS"), 'H');
assert_eq!(amino_acid_letter("TRP"), 'W');
assert_eq!(amino_acid_letter("UNK"), 'X');
}
#[test]
fn interactive_frames_use_native_terminal_text() {
let protein = Protein::from_pdb_str(PICK_PDB, "pick.pdb").unwrap();
let viewer = Viewer::new(protein, Representation::Backbone, ColorScheme::Confidence);
let interactive = viewer.render_scene(100, 72, false, Some((100, 36)));
assert!(
interactive
.text
.iter()
.any(|run| run.text.contains("SEQUENCE"))
);
assert!(
interactive
.text
.iter()
.any(|run| run.text.contains("hover inspect"))
);
let snapshot = viewer.snapshot(800, 600);
assert!(snapshot.text.is_empty());
}
#[test]
fn hover_card_border_encloses_native_text() {
let protein = Protein::from_pdb_str(PICK_PDB, "pick.pdb").unwrap();
let mut viewer = Viewer::new(protein, Representation::Backbone, ColorScheme::Confidence);
viewer.hovered_residue = Some(0);
let surface = TerminalSurface {
width: 100,
height: 72,
columns: 100,
rows: 36,
};
let layout = viewer.terminal_view_layout(
surface.width,
surface.height,
surface.columns,
surface.rows,
);
let card = terminal_card_layout(layout, surface);
let frame = viewer.render_scene(
surface.width,
surface.height,
false,
Some((surface.columns, surface.rows)),
);
let hover = frame
.text
.iter()
.find(|run| run.text.starts_with("HOVER"))
.expect("hover card text");
assert!(hover.column > card.column);
assert!(
hover.column + (hover.text.len() as u16)
< card.column.saturating_add(card.width_columns)
);
assert!(hover.row >= card.row && hover.row < card.row + 3);
let border = [78, 103, 132, 255];
for (x, y) in [
(card.rect.x, card.rect.y),
(
card.rect.x + card.rect.width - 1,
card.rect.y + card.rect.height - 1,
),
] {
let offset = ((y * frame.width + x) * 4) as usize;
assert_eq!(&frame.pixels[offset..offset + 4], &border);
}
}
}