lightdock/
swarm.rs

1use super::glowworm::distance;
2use super::glowworm::Glowworm;
3use super::qt::Quaternion;
4use super::scoring::Score;
5use rand::Rng;
6use std::fs::File;
7use std::io::{Error, Write};
8
9pub struct Swarm<'a> {
10    pub glowworms: Vec<Glowworm<'a>>,
11}
12
13impl<'a> Default for Swarm<'a> {
14    fn default() -> Self {
15        Swarm::new()
16    }
17}
18
19impl<'a> Swarm<'a> {
20    pub fn new() -> Self {
21        Swarm {
22            glowworms: Vec::new(),
23        }
24    }
25
26    pub fn add_glowworms(
27        &mut self,
28        positions: &[Vec<f64>],
29        scoring: &'a Box<dyn Score>,
30        use_anm: bool,
31        rec_num_anm: usize,
32        lig_num_anm: usize,
33    ) {
34        for (i, position) in positions.iter().enumerate() {
35            // Translation component
36            let translation = vec![position[0], position[1], position[2]];
37            // Rotation component
38            let rotation = Quaternion::new(position[3], position[4], position[5], position[6]);
39            // ANM for receptor
40            let mut rec_nmodes: Vec<f64> = Vec::new();
41            if use_anm && rec_num_anm > 0 {
42                for j in 7..7 + rec_num_anm {
43                    rec_nmodes.push(positions[i][j]);
44                }
45            }
46            // ANM for ligand
47            let mut lig_nmodes: Vec<f64> = Vec::new();
48            if use_anm && lig_num_anm > 0 {
49                for j in 7 + rec_num_anm..positions[i].len() {
50                    lig_nmodes.push(positions[i][j]);
51                }
52            }
53            let glowworm = Glowworm::new(
54                i as u32,
55                translation,
56                rotation,
57                rec_nmodes,
58                lig_nmodes,
59                scoring,
60                use_anm,
61            );
62            self.glowworms.push(glowworm);
63        }
64    }
65
66    pub fn update_luciferin(&mut self) {
67        for glowworm in self.glowworms.iter_mut() {
68            glowworm.compute_luciferin();
69        }
70    }
71
72    pub fn movement_phase(&mut self, rng: &mut rand::prelude::StdRng) {
73        // Save original positions
74        let mut positions: Vec<Vec<f64>> = Vec::new();
75        let mut rotations: Vec<Quaternion> = Vec::new();
76        let mut anm_recs: Vec<Vec<f64>> = Vec::new();
77        let mut anm_ligs: Vec<Vec<f64>> = Vec::new();
78        for glowworm in self.glowworms.iter() {
79            positions.push(glowworm.translation.clone());
80            rotations.push(glowworm.rotation);
81            anm_recs.push(glowworm.rec_nmodes.clone());
82            anm_ligs.push(glowworm.lig_nmodes.clone());
83        }
84
85        // First search for each glowworm's neighbors
86        let mut neighbors: Vec<Vec<u32>> = Vec::new();
87        for i in 0..self.glowworms.len() {
88            let mut this_neighbors = Vec::new();
89            let g1 = &self.glowworms[i];
90            for j in 0..self.glowworms.len() {
91                if i != j {
92                    let g2 = &self.glowworms[j];
93                    if g1.luciferin < g2.luciferin {
94                        let distance = distance(g1, g2);
95                        if distance < g1.vision_range {
96                            this_neighbors.push(g2.id);
97                        }
98                    }
99                }
100            }
101            neighbors.push(this_neighbors);
102        }
103
104        // Second compute probability moving towards the neighbor
105        let mut luciferins = Vec::new();
106        for glowworm in self.glowworms.iter_mut() {
107            luciferins.push(glowworm.luciferin);
108        }
109        for i in 0..self.glowworms.len() {
110            let glowworm = &mut self.glowworms[i];
111            glowworm.neighbors = neighbors[i].clone();
112            glowworm.compute_probability_moving_toward_neighbor(&luciferins);
113        }
114
115        // Finally move to the selected position
116        for i in 0..self.glowworms.len() {
117            let glowworm = &mut self.glowworms[i];
118            let neighbor_id = glowworm.select_random_neighbor(rng.gen::<f64>());
119            let position = &positions[neighbor_id as usize];
120            let rotation = &rotations[neighbor_id as usize];
121            let anm_rec = &anm_recs[neighbor_id as usize];
122            let anm_lig = &anm_ligs[neighbor_id as usize];
123            glowworm.move_towards(neighbor_id, position, rotation, anm_rec, anm_lig);
124            glowworm.update_vision_range();
125        }
126    }
127
128    pub fn save(&mut self, step: u32, output_directory: &str) -> Result<(), Error> {
129        let path = format!("{}/gso_{:?}.out", output_directory, step);
130        let mut output = File::create(path)?;
131        writeln!(
132            output,
133            "#Coordinates  RecID  LigID  Luciferin  Neighbor's number  Vision Range  Scoring"
134        )?;
135        for glowworm in self.glowworms.iter() {
136            write!(
137                output,
138                "({:.7}, {:.7}, {:.7}, {:.7}, {:.7}, {:.7}, {:.7}",
139                glowworm.translation[0],
140                glowworm.translation[1],
141                glowworm.translation[2],
142                glowworm.rotation.w,
143                glowworm.rotation.x,
144                glowworm.rotation.y,
145                glowworm.rotation.z
146            )?;
147            if glowworm.use_anm && !glowworm.rec_nmodes.is_empty() {
148                for i in 0..glowworm.rec_nmodes.len() {
149                    write!(output, ", {:.7}", glowworm.rec_nmodes[i])?;
150                }
151            }
152            if glowworm.use_anm && !glowworm.lig_nmodes.is_empty() {
153                for i in 0..glowworm.lig_nmodes.len() {
154                    write!(output, ", {:.7}", glowworm.lig_nmodes[i])?;
155                }
156            }
157            writeln!(
158                output,
159                ")    0    0   {:.8}  {:?} {:.3} {:.8}",
160                glowworm.luciferin,
161                glowworm.neighbors.len(),
162                glowworm.vision_range,
163                glowworm.scoring
164            )?;
165        }
166        Ok(())
167    }
168}