Skip to main content

pfx_physics/
flock.rs

1use crate::math::{add3, clamp3, len3, limit3, norm3, scale3, sub3};
2use crate::noise::Rng;
3
4#[derive(Clone, Copy, Debug)]
5pub struct Boid {
6    pub pos: [f32; 3],
7    pub vel: [f32; 3],
8}
9
10#[derive(Clone, Copy, Debug)]
11pub struct FlockParams {
12    pub separation: f32,
13    pub separation_weight: f32,
14    pub alignment: f32,
15    pub alignment_weight: f32,
16    pub cohesion: f32,
17    pub cohesion_weight: f32,
18    pub max_speed: f32,
19    pub max_force: f32,
20    pub bounds_min: [f32; 3],
21    pub bounds_max: [f32; 3],
22    pub bounds_weight: f32,
23    pub margin: f32,
24    pub target: Option<[f32; 3]>,
25    pub target_weight: f32,
26}
27
28impl FlockParams {
29    pub fn birds() -> Self {
30        Self {
31            separation: 1.6,
32            separation_weight: 1.8,
33            alignment: 4.0,
34            alignment_weight: 0.8,
35            cohesion: 6.0,
36            cohesion_weight: 0.45,
37            max_speed: 6.0,
38            max_force: 14.0,
39            bounds_min: [-24.0, 1.0, -24.0],
40            bounds_max: [24.0, 14.0, 24.0],
41            bounds_weight: 3.0,
42            margin: 2.0,
43            target: None,
44            target_weight: 1.2,
45        }
46    }
47
48    pub fn fish() -> Self {
49        Self {
50            separation: 0.8,
51            separation_weight: 2.2,
52            alignment: 2.2,
53            alignment_weight: 1.1,
54            cohesion: 3.0,
55            cohesion_weight: 0.7,
56            max_speed: 3.5,
57            max_force: 10.0,
58            bounds_min: [-16.0, -0.6, -16.0],
59            bounds_max: [16.0, 0.6, 16.0],
60            bounds_weight: 4.0,
61            margin: 0.8,
62            target: None,
63            target_weight: 1.4,
64        }
65    }
66}
67
68pub struct Flock {
69    boids: Vec<Boid>,
70    params: FlockParams,
71}
72
73impl Flock {
74    pub fn new(boids: Vec<Boid>, params: FlockParams) -> Self {
75        Self { boids, params }
76    }
77
78    pub fn seeded(count: usize, seed: u64, params: FlockParams) -> Self {
79        let mut rng = Rng::new(seed);
80        let mut boids = Vec::with_capacity(count);
81        for _ in 0..count {
82            let pos = [0, 1, 2].map(|k| rng.range(params.bounds_min[k], params.bounds_max[k]));
83            let vel = norm3([
84                rng.range(-1.0, 1.0),
85                rng.range(-1.0, 1.0),
86                rng.range(-1.0, 1.0),
87            ]);
88            boids.push(Boid {
89                pos,
90                vel: scale3(vel, params.max_speed * 0.5),
91            });
92        }
93        Self { boids, params }
94    }
95
96    pub fn boids(&self) -> &[Boid] {
97        &self.boids
98    }
99
100    pub fn follow(&mut self, target: Option<[f32; 3]>) {
101        self.params.target = target;
102    }
103
104    pub fn step(&mut self, dt: f32) {
105        if !dt.is_finite() || dt <= 0.0 || self.boids.is_empty() {
106            return;
107        }
108        let n = self.boids.len();
109        let params = self.params;
110        let mut acc = vec![[0.0; 3]; n];
111        for (i, slot) in acc.iter_mut().enumerate() {
112            let mut sep = [0.0; 3];
113            let mut align = [0.0; 3];
114            let mut centre = [0.0; 3];
115            let mut align_n = 0.0;
116            let mut cohere_n = 0.0;
117            for j in 0..n {
118                if i == j {
119                    continue;
120                }
121                let d = sub3(self.boids[i].pos, self.boids[j].pos);
122                let dist = len3(d).max(1e-4);
123                if dist < params.separation {
124                    let push = scale3(norm3(d), (params.separation - dist) / params.separation);
125                    sep = add3(sep, push);
126                }
127                if dist < params.alignment {
128                    align = add3(align, self.boids[j].vel);
129                    align_n += 1.0;
130                }
131                if dist < params.cohesion {
132                    centre = add3(centre, self.boids[j].pos);
133                    cohere_n += 1.0;
134                }
135            }
136            let mut a = scale3(sep, params.separation_weight);
137            if align_n > 0.0 {
138                let want = scale3(norm3(scale3(align, 1.0 / align_n)), params.max_speed);
139                a = add3(
140                    a,
141                    scale3(sub3(want, self.boids[i].vel), params.alignment_weight),
142                );
143            }
144            if cohere_n > 0.0 {
145                let want = scale3(
146                    norm3(sub3(scale3(centre, 1.0 / cohere_n), self.boids[i].pos)),
147                    params.max_speed,
148                );
149                a = add3(
150                    a,
151                    scale3(sub3(want, self.boids[i].vel), params.cohesion_weight),
152                );
153            }
154            a = add3(a, bounds_force(self.boids[i].pos, &params));
155            if let Some(target) = params.target {
156                let want = scale3(norm3(sub3(target, self.boids[i].pos)), params.max_speed);
157                a = add3(
158                    a,
159                    scale3(sub3(want, self.boids[i].vel), params.target_weight),
160                );
161            }
162            *slot = limit3(a, params.max_force);
163        }
164        for (boid, a) in self.boids.iter_mut().zip(acc) {
165            boid.vel = limit3(add3(boid.vel, scale3(a, dt)), params.max_speed);
166            boid.pos = clamp3(
167                add3(boid.pos, scale3(boid.vel, dt)),
168                params.bounds_min,
169                params.bounds_max,
170            );
171        }
172    }
173}
174
175fn bounds_force(p: [f32; 3], params: &FlockParams) -> [f32; 3] {
176    let mut a = [0.0; 3];
177    for k in 0..3 {
178        let lo = params.bounds_min[k] + params.margin;
179        let hi = params.bounds_max[k] - params.margin;
180        if p[k] < lo {
181            a[k] += params.bounds_weight * (lo - p[k]);
182        } else if p[k] > hi {
183            a[k] -= params.bounds_weight * (p[k] - hi);
184        }
185    }
186    a
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    fn apart(a: [f32; 3], b: [f32; 3]) -> f32 {
194        len3(sub3(a, b))
195    }
196
197    #[test]
198    fn separation_pushes_a_close_pair_apart() {
199        let mut params = FlockParams::birds();
200        params.separation = 2.0;
201        params.separation_weight = 12.0;
202        params.alignment_weight = 0.0;
203        params.cohesion_weight = 0.0;
204        params.target = None;
205        params.bounds_min = [-50.0; 3];
206        params.bounds_max = [50.0; 3];
207        params.max_speed = 6.0;
208        params.max_force = 40.0;
209        let mut flock = Flock::new(
210            vec![
211                Boid {
212                    pos: [0.0, 0.0, 0.0],
213                    vel: [0.0; 3],
214                },
215                Boid {
216                    pos: [0.2, 0.0, 0.0],
217                    vel: [0.0; 3],
218                },
219            ],
220            params,
221        );
222        let start = apart(flock.boids()[0].pos, flock.boids()[1].pos);
223        for _ in 0..90 {
224            flock.step(1.0 / 60.0);
225        }
226        let now = apart(flock.boids()[0].pos, flock.boids()[1].pos);
227        assert!(now > start, "{now} vs {start}");
228        assert!(now > 1.0, "{now}");
229    }
230
231    #[test]
232    fn the_same_seed_repeats_the_flock() {
233        let params = FlockParams::fish();
234        let mut a = Flock::seeded(80, 42, params);
235        let mut b = Flock::seeded(80, 42, params);
236        a.follow(Some([2.0, 0.0, -1.0]));
237        b.follow(Some([2.0, 0.0, -1.0]));
238        for _ in 0..120 {
239            a.step(1.0 / 60.0);
240            b.step(1.0 / 60.0);
241        }
242        for (left, right) in a.boids().iter().zip(b.boids()) {
243            assert_eq!(left.pos, right.pos);
244            assert_eq!(left.vel, right.vel);
245        }
246        assert!(a.boids().iter().all(|boid| {
247            (0..3)
248                .all(|k| boid.pos[k] >= params.bounds_min[k] && boid.pos[k] <= params.bounds_max[k])
249        }));
250    }
251
252    #[test]
253    fn a_flock_follows_its_target_and_stays_in_bounds() {
254        let mut params = FlockParams::birds();
255        params.target = Some([10.0, 6.0, 0.0]);
256        let mut flock = Flock::seeded(200, 7, params);
257        let start = flock.boids().iter().map(|b| b.pos[0]).sum::<f32>() / 200.0;
258        for _ in 0..180 {
259            flock.step(1.0 / 60.0);
260        }
261        let end = flock.boids().iter().map(|b| b.pos[0]).sum::<f32>() / 200.0;
262        assert!(end > start, "{end} vs {start}");
263        assert!(flock.boids().iter().all(|boid| {
264            boid.pos.iter().all(|v| v.is_finite())
265                && (0..3).all(|k| {
266                    boid.pos[k] >= params.bounds_min[k] && boid.pos[k] <= params.bounds_max[k]
267                })
268        }));
269    }
270}