1use crate::model::structure::Structure;
2use crate::model::types::Point;
3use nalgebra::{Rotation3, Vector3};
4
5pub struct Transform;
6
7impl Transform {
8 pub fn translate(structure: &mut Structure, x: f64, y: f64, z: f64) {
9 let translation = Vector3::new(x, y, z);
10 for atom in structure.iter_atoms_mut() {
11 atom.translate_by(&translation);
12 }
13 }
14
15 pub fn center_geometry(structure: &mut Structure, target: Option<Point>) {
16 let current_center = structure.geometric_center();
17 let target_point = target.unwrap_or(Point::origin());
18 let translation = target_point - current_center;
19
20 for atom in structure.iter_atoms_mut() {
21 atom.translate_by(&translation);
22 }
23 }
24
25 pub fn center_mass(structure: &mut Structure, target: Option<Point>) {
26 let current_com = structure.center_of_mass();
27 let target_point = target.unwrap_or(Point::origin());
28 let translation = target_point - current_com;
29
30 for atom in structure.iter_atoms_mut() {
31 atom.translate_by(&translation);
32 }
33 }
34
35 pub fn rotate_x(structure: &mut Structure, radians: f64) {
36 let rotation = Rotation3::from_axis_angle(&Vector3::x_axis(), radians);
37 Self::apply_rotation(structure, rotation);
38 }
39
40 pub fn rotate_y(structure: &mut Structure, radians: f64) {
41 let rotation = Rotation3::from_axis_angle(&Vector3::y_axis(), radians);
42 Self::apply_rotation(structure, rotation);
43 }
44
45 pub fn rotate_z(structure: &mut Structure, radians: f64) {
46 let rotation = Rotation3::from_axis_angle(&Vector3::z_axis(), radians);
47 Self::apply_rotation(structure, rotation);
48 }
49
50 pub fn rotate_euler(structure: &mut Structure, x_rad: f64, y_rad: f64, z_rad: f64) {
51 let rotation = Rotation3::from_euler_angles(x_rad, y_rad, z_rad);
52 Self::apply_rotation(structure, rotation);
53 }
54
55 fn apply_rotation(structure: &mut Structure, rotation: Rotation3<f64>) {
56 for atom in structure.iter_atoms_mut() {
57 atom.pos = rotation * atom.pos;
58 }
59
60 if let Some(box_vecs) = structure.box_vectors {
61 let v1 = Vector3::from(box_vecs[0]);
62 let v2 = Vector3::from(box_vecs[1]);
63 let v3 = Vector3::from(box_vecs[2]);
64
65 let v1_rot = rotation * v1;
66 let v2_rot = rotation * v2;
67 let v3_rot = rotation * v3;
68
69 structure.box_vectors = Some([v1_rot.into(), v2_rot.into(), v3_rot.into()]);
70 }
71 }
72}
73
74#[cfg(test)]
75mod tests {
76 use super::Transform;
77 use crate::model::{
78 atom::Atom,
79 chain::Chain,
80 residue::Residue,
81 structure::Structure,
82 types::{Element, Point, ResidueCategory, StandardResidue},
83 };
84
85 fn structure_with_points(points: &[Point]) -> Structure {
86 let mut chain = Chain::new("A");
87 let mut residue = Residue::new(
88 1,
89 None,
90 "GLY",
91 Some(StandardResidue::GLY),
92 ResidueCategory::Standard,
93 );
94
95 for (idx, point) in points.iter().enumerate() {
96 let name = format!("C{}", idx);
97 residue.add_atom(Atom::new(&name, Element::C, *point));
98 }
99
100 chain.add_residue(residue);
101 let mut structure = Structure::new();
102 structure.add_chain(chain);
103 structure
104 }
105
106 fn assert_point_close(actual: &Point, expected: &Point) {
107 assert!((actual.x - expected.x).abs() < 1e-6);
108 assert!((actual.y - expected.y).abs() < 1e-6);
109 assert!((actual.z - expected.z).abs() < 1e-6);
110 }
111
112 #[test]
113 fn translate_moves_all_atoms_by_vector() {
114 let mut structure =
115 structure_with_points(&[Point::new(0.0, 0.0, 0.0), Point::new(1.0, 2.0, 3.0)]);
116
117 Transform::translate(&mut structure, 5.0, -2.0, 1.5);
118
119 let mut atoms = structure.iter_atoms();
120 assert_point_close(&atoms.next().unwrap().pos, &Point::new(5.0, -2.0, 1.5));
121 assert_point_close(&atoms.next().unwrap().pos, &Point::new(6.0, 0.0, 4.5));
122 }
123
124 #[test]
125 fn center_geometry_moves_geometric_center_to_target() {
126 let mut structure =
127 structure_with_points(&[Point::new(2.0, 0.0, 0.0), Point::new(4.0, 0.0, 0.0)]);
128
129 Transform::center_geometry(&mut structure, Some(Point::new(10.0, 0.0, 0.0)));
130
131 let center = structure.geometric_center();
132 assert_point_close(¢er, &Point::new(10.0, 0.0, 0.0));
133 }
134
135 #[test]
136 fn center_mass_moves_center_of_mass_to_origin_by_default() {
137 let mut structure =
138 structure_with_points(&[Point::new(2.0, 0.0, 0.0), Point::new(4.0, 0.0, 0.0)]);
139
140 {
141 let chain = structure.chain_mut("A").unwrap();
142 let residue = chain.residue_mut(1, None).unwrap();
143 residue
144 .iter_atoms_mut()
145 .enumerate()
146 .for_each(|(idx, atom)| {
147 atom.element = if idx == 0 { Element::H } else { Element::O };
148 });
149 }
150
151 Transform::center_mass(&mut structure, None);
152
153 let com = structure.center_of_mass();
154 assert_point_close(&com, &Point::origin());
155 }
156
157 #[test]
158 fn rotate_z_rotates_atoms_about_origin() {
159 let mut structure =
160 structure_with_points(&[Point::new(1.0, 0.0, 0.0), Point::new(0.0, 2.0, 0.0)]);
161
162 Transform::rotate_z(&mut structure, std::f64::consts::FRAC_PI_2);
163
164 let mut atoms = structure.iter_atoms();
165 assert_point_close(&atoms.next().unwrap().pos, &Point::new(0.0, 1.0, 0.0));
166 assert_point_close(&atoms.next().unwrap().pos, &Point::new(-2.0, 0.0, 0.0));
167 }
168
169 #[test]
170 fn rotate_euler_updates_box_vectors() {
171 let mut structure = structure_with_points(&[Point::new(1.0, 0.0, 0.0)]);
172 structure.box_vectors = Some([[1.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 3.0]]);
173
174 Transform::rotate_euler(&mut structure, 0.0, 0.0, std::f64::consts::FRAC_PI_2);
175
176 let box_vectors = structure.box_vectors.unwrap();
177 assert_point_close(&Point::from(box_vectors[0]), &Point::new(0.0, 1.0, 0.0));
178 assert_point_close(&Point::from(box_vectors[1]), &Point::new(-2.0, 0.0, 0.0));
179 assert_point_close(&Point::from(box_vectors[2]), &Point::new(0.0, 0.0, 3.0));
180 }
181}