pdbvis 0.1.2

A 3D protein structure viewer that loads and visualizes proteins from the Protein Data Bank (PDB)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213

use crate::protein::Protein;
use anyhow::Result;
use kiss3d::camera::{ArcBall, Camera};
use kiss3d::event::{Action, Key};
use kiss3d::nalgebra::{Point3, Translation3, UnitQuaternion, Vector3};
use kiss3d::scene::SceneNode;
use kiss3d::window::Window;

#[derive(PartialEq)]
enum ViewMode {
    Backbone,
    Cartoon,
}

pub struct ProteinVisualizer<'a> {
    window: &'a mut Window,
    camera: ArcBall,
    protein: &'a Protein,
    view_mode: ViewMode,
    backbone_group: SceneNode,
    cartoon_group: Option<SceneNode>,
    key_pressed: bool,
    axes: [SceneNode; 3], // X, Y, Z axes
}

impl<'a> ProteinVisualizer<'a> {
    pub fn new(window: &'a mut Window, protein: &'a Protein) -> Result<Self> {
        // Create both visualizations immediately
        let mut backbone_group = window.add_group();
        create_backbone_visualization(protein, &mut backbone_group);
        backbone_group.set_visible(false); // Hide backbone initially

        let mut cartoon_group = window.add_group();
        create_cartoon_visualization(protein, &mut cartoon_group);

        // Setup camera with better initial position
        let camera_distance = protein.maximum_radius * 3.0;
        let eye = Point3::new(camera_distance, camera_distance * 0.5, camera_distance);
        let mut camera = ArcBall::new(eye, Point3::origin());
        camera.set_up_axis(Vector3::new(0.0, 1.0, 0.0));

        // Set minimum distance (maximum zoom)
        camera.set_min_dist(protein.maximum_radius * 1.0);

        // Set maximum distance (minimum zoom)
        camera.set_max_dist(protein.maximum_radius * 5.0);

        // Create coordinate axes
        let axes = create_axes(window, protein.maximum_radius * 0.5);

        Ok(Self {
            window,
            camera,
            protein,
            view_mode: ViewMode::Cartoon,
            backbone_group,
            cartoon_group: Some(cartoon_group),
            key_pressed: false,
            axes,
        })
    }

    pub fn run(&mut self) -> Result<()> {
        while self.window.render_with_camera(&mut self.camera) {
            // Handle view mode switching with semicolon key
            if self.window.get_key(Key::Semicolon) == Action::Press {
                if !self.key_pressed {
                    self.key_pressed = true;
                    self.toggle_view_mode();
                }
            } else {
                self.key_pressed = false;
            }

            // Update axes positions to stay in view corner
            self.update_axes_position();
        }
        Ok(())
    }

    fn update_axes_position(&mut self) {
        let view = self.camera.view_transform();
        let scale = self.protein.maximum_radius * 0.15;

        let offset = Vector3::new(
            self.window.size()[0] as f32 * 0.4,
            -(self.window.size()[1] as f32) * 0.4,
            0.0,
        );

        for (_i, axis) in self.axes.iter_mut().enumerate() {
            axis.set_local_scale(scale, scale, scale);
            axis.set_local_rotation(view.rotation);
            axis.set_local_translation(Translation3::from(offset));
        }
    }

    fn toggle_view_mode(&mut self) {
        match self.view_mode {
            ViewMode::Backbone => {
                self.view_mode = ViewMode::Cartoon;
                self.backbone_group.set_visible(false);
                if let Some(ref mut group) = self.cartoon_group {
                    group.set_visible(true);
                }
            }
            ViewMode::Cartoon => {
                self.view_mode = ViewMode::Backbone;
                self.backbone_group.set_visible(true);
                if let Some(ref mut group) = self.cartoon_group {
                    group.set_visible(false);
                }
            }
        }
    }
}

fn create_axes(window: &mut Window, _size: f32) -> [SceneNode; 3] {
    let mut x_axis = window.add_cylinder(0.1, 1.0);
    let mut y_axis = window.add_cylinder(0.1, 1.0);
    let mut z_axis = window.add_cylinder(0.1, 1.0);

    // Set colors for each axis (RGB)
    x_axis.set_color(1.0, 0.0, 0.0); // Red for X
    y_axis.set_color(0.0, 1.0, 0.0); // Green for Y
    z_axis.set_color(0.0, 0.0, 1.0); // Blue for Z

    // Set rotations for Y and Z axes
    y_axis.set_local_rotation(UnitQuaternion::from_axis_angle(
        &Vector3::x_axis(),
        std::f32::consts::FRAC_PI_2,
    ));
    z_axis.set_local_rotation(UnitQuaternion::from_axis_angle(
        &Vector3::y_axis(),
        std::f32::consts::FRAC_PI_2,
    ));

    [x_axis, y_axis, z_axis]
}

fn create_backbone_visualization(protein: &Protein, group: &mut SceneNode) {
    let mut previous_point = None;
    for atom in &protein.atoms {
        if atom.atom_type == "CA" {
            let centered_pos = atom.position - protein.center.coords;
            let mut sphere = group.add_sphere(0.5);

            // Color spheres based on temperature factor (B-factor)
            let normalized_temp = atom.temperature_factor / 100.0; // Typical B-factor range 0-100
            sphere.set_color(
                normalized_temp,               // More red = higher temperature factor
                0.7 * (1.0 - normalized_temp), // Less blue = higher temperature factor
                1.0 - normalized_temp,         // Less green = higher temperature factor
            );

            // Scale sphere based on occupancy
            let scale = atom.occupancy * 0.5;
            sphere.set_local_scale(scale, scale, scale);

            sphere.set_local_translation(Translation3::from(centered_pos));

            if let Some(previous) = previous_point {
                let mut line = group.add_cylinder(0.1, 1.0);
                // Color lines based on chain ID
                match atom.chain_id {
                    'A' => line.set_color(1.0, 0.0, 0.0), // Red for chain A
                    'B' => line.set_color(0.0, 1.0, 0.0), // Green for chain B
                    _ => line.set_color(0.0, 0.0, 1.0),   // Blue for other chains
                }

                let direction = centered_pos - previous;
                let middle_point = previous + direction * 0.5;
                line.set_local_translation(Translation3::from(middle_point));
            }
            previous_point = Some(centered_pos);
        }
    }
}

fn create_cartoon_visualization(protein: &Protein, group: &mut SceneNode) {
    let mut previous_point = None;
    let mut previous_direction = None;

    for atom in &protein.atoms {
        if atom.atom_type == "CA" {
            let centered_pos = atom.position - protein.center.coords;

            if let Some(previous) = previous_point {
                let direction: Vector3<f32> = centered_pos - previous;
                let length = direction.magnitude();

                // Create ribbon segment (thinner and taller)
                let mut ribbon = group.add_cube(0.4, 0.8, length);
                ribbon.set_color(0.0, 0.8, 0.2);

                // Position at midpoint
                let middle_point = previous + direction * 0.5;
                ribbon.set_local_translation(Translation3::from(middle_point));

                // Orient ribbon
                if let Some(_previous_direction) = previous_direction {
                    let rotation =
                        UnitQuaternion::rotation_between(&Vector3::z(), &direction.normalize())
                            .unwrap_or(UnitQuaternion::identity());
                    ribbon.set_local_rotation(rotation);
                }

                previous_direction = Some(direction);
            }
            previous_point = Some(centered_pos);
        }
    }
}