use std::sync::Arc;
use rstui::{App, AppContext, AppEvent, ControlFlow, Frame};
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,
}
}
}
pub struct Viewer {
protein: Arc<Protein>,
yaw: f32,
pitch: f32,
zoom: f32,
representation: Representation,
color_scheme: ColorScheme,
dragging: Option<(u32, u32)>,
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 {
Self {
protein: Arc::new(protein),
yaw: -0.72,
pitch: 0.38,
zoom: 1.35,
representation,
color_scheme,
dragging: 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
}
fn reset_view(&mut self) {
self.yaw = -0.72;
self.pitch = 0.38;
self.zoom = 1.35;
self.auto_rotate = false;
}
pub fn snapshot(&self, width: u32, height: u32) -> Frame {
self.render_scene(width.max(1), height.max(1), false)
}
fn render_scene(&self, width: u32, height: u32, fast: bool) -> Frame {
let mut canvas = Canvas::new(width, height, BACKGROUND_TOP);
canvas.vertical_gradient(BACKGROUND_TOP, BACKGROUND_BOTTOM);
let detailed = width >= 280 && height >= 105;
let header = if detailed { 22 } else { 2 };
let footer = if detailed { 22 } else { 2 };
let scene_height = height.saturating_sub(header + footer).max(1);
let molecule_width = if self.atlas.is_some() && width >= 360 {
(width as f32 * 0.62).round() as u32
} else {
width
}
.max(1);
let scene_center = [
molecule_width as f32 * 0.5,
header as f32 + scene_height as f32 * 0.5,
];
let scene_scale = molecule_width.min(scene_height) as f32 * 0.43 * self.zoom;
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(atlas) = &self.atlas
&& molecule_width < width
{
let panel = Rect {
x: molecule_width,
y: header,
width: width - molecule_width,
height: scene_height,
};
render_pae_panel(&mut canvas, panel, atlas);
}
if detailed {
canvas.fill_rect(0, 0, width, 20, [3, 7, 13, 215]);
canvas.fill_rect(0, height.saturating_sub(20), width, 20, [3, 7, 13, 225]);
canvas.horizontal_line(20, [35, 55, 71, 255]);
canvas.horizontal_line(height.saturating_sub(21), [35, 55, 71, 255]);
let title_capacity = molecule_width.saturating_sub(16) as usize / 8;
canvas.text(8, 6, &truncate(&self.protein.title, title_capacity), TEXT);
let status = format!(
"{} atoms {} residues {} chains mode {} color {}{}",
self.protein.atoms.len(),
self.protein.residue_count,
self.protein.chains.len(),
self.representation.name(),
self.color_scheme.name(),
if self.auto_rotate { " AUTO" } else { "" }
);
canvas.text(8, height.saturating_sub(16), &status, ACCENT);
let help = "drag/arrows rotate wheel +/- zoom m mode c color space auto r reset q quit";
let help_x = width.saturating_sub(help.len() as u32 * 8 + 8);
if help_x > status.len() as u32 * 8 + 24 {
canvas.text(help_x, height.saturating_sub(16), help, MUTED);
}
}
Frame::rgba(width, height, canvas.pixels).expect("canvas dimensions are valid")
}
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.protein.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 {
let atom = &self.protein.atoms[sample.atom_index];
let local = (sample.center - self.protein.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
}
_ => false,
},
AppEvent::Mouse {
kind,
column,
row,
pixel_x,
pixel_y,
button,
} => {
let x = pixel_x.unwrap_or(u32::from(column) * 8);
let y = pixel_y.unwrap_or(u32::from(row) * 16);
match (kind.as_str(), button.as_deref()) {
("down", Some("left")) => {
self.dragging = Some((x, y));
false
}
("drag", Some("left")) => {
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;
false
}
("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 scale = if context.quality == "fast"
&& matches!(context.protocol.as_str(), "kitty" | "iterm2")
{
0.58
} else {
1.0
};
let width = ((context.geometry.pixel_width as f32 * scale).round() as u32).max(1);
let height = ((context.geometry.pixel_height as f32 * scale).round() as u32).max(1);
Ok(self.render_scene(width, height, context.quality == "fast"))
}
}
#[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,
}
#[derive(Clone, Copy)]
struct RibbonNode {
center: Vec3,
normal: Vec3,
half_width: f32,
half_thickness: f32,
secondary: SecondaryStructure,
atom_index: usize,
}
#[derive(Clone, Copy)]
struct Rect {
x: u32,
y: u32,
width: u32,
height: u32,
}
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,
});
}
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,
}];
}
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
},
});
}
}
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,
});
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) {
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";
canvas.text(
panel.x + 10,
panel.y + 8,
&truncate(title, panel.width.saturating_sub(20) as usize / 8),
TEXT,
);
let Some(confidence) = &atlas.confidence else {
canvas.text(panel.x + 10, panel.y + 38, &atlas.label, ACCENT);
canvas.text(panel.x + 10, panel.y + 58, "PAE is unavailable for", MUTED);
canvas.text(
panel.x + 10,
panel.y + 70,
"live ESMFold predictions.",
MUTED,
);
return;
};
let available_height = panel.height.saturating_sub(104);
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 + 32;
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);
canvas.text(heat_x, heat_y + heat_size + 4, "0", MUTED);
let end_label = last.to_string();
canvas.text(
heat_x.saturating_add(heat_size.saturating_sub(end_label.len() as u32 * 8)),
heat_y + heat_size + 4,
&end_label,
MUTED,
);
let legend_y = heat_y + heat_size + 20;
let legend_width = heat_size.min(150);
let legend_x = panel.x + (panel.width.saturating_sub(legend_width)) / 2;
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, 8, pae_color(value));
}
canvas.text(legend_x, legend_y + 12, "0", TEXT);
canvas.text(legend_x + legend_width / 2 - 8, legend_y + 12, "15", TEXT);
canvas.text(
legend_x + legend_width.saturating_sub(16),
legend_y + 12,
"30",
TEXT,
);
let metric = format!("{} pTM {:.3}", atlas.label, confidence.ptm);
canvas.text(
panel.x + 10,
panel.y + panel.height.saturating_sub(14),
&truncate(&metric, panel.width.saturating_sub(20) as usize / 8),
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 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 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::*;
#[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);
}
}