docs/bevy/examples/ecs/iter_combinations.rs 5.1 K · 159 lines · raw · history

1 //! Shows how to iterate over combinations of query results.
2
3 use bevy::{color::palettes::css::ORANGE_RED, math::FloatPow, prelude::*};
4 use chacha20::ChaCha8Rng;
5 use rand::{RngExt, SeedableRng};
6
7 fn main() {
8 App::new()
9 .add_plugins(DefaultPlugins)
10 .insert_resource(ClearColor(Color::BLACK))
11 .add_systems(Startup, generate_bodies)
12 .add_systems(FixedUpdate, (interact_bodies, integrate))
13 .add_systems(Update, look_at_star)
14 .run();
15 }
16
17 const GRAVITY_CONSTANT: f32 = 0.001;
18 const NUM_BODIES: usize = 100;
19
20 #[derive(Component, Default)]
21 struct Mass(f32);
22 #[derive(Component, Default)]
23 struct Acceleration(Vec3);
24 #[derive(Component, Default)]
25 struct LastPos(Vec3);
26 #[derive(Component)]
27 struct Star;
28
29 #[derive(Bundle, Default)]
30 struct BodyBundle {
31 mesh: Mesh3d,
32 material: MeshMaterial3d<StandardMaterial>,
33 mass: Mass,
34 last_pos: LastPos,
35 acceleration: Acceleration,
36 }
37
38 fn generate_bodies(
39 time: Res<Time<Fixed>>,
40 mut commands: Commands,
41 mut meshes: ResMut<Assets<Mesh>>,
42 mut materials: ResMut<Assets<StandardMaterial>>,
43 ) {
44 let mesh = meshes.add(Sphere::new(1.0).mesh().ico(3).unwrap());
45
46 let color_range = 0.5..1.0;
47 let vel_range = -0.5..0.5;
48
49 // We're seeding the PRNG here to make this example deterministic for testing purposes.
50 // This isn't strictly required in practical use unless you need your app to be deterministic.
51 let mut rng = ChaCha8Rng::seed_from_u64(19878367467713);
52 for _ in 0..NUM_BODIES {
53 let radius: f32 = rng.random_range(0.1..0.7);
54 let mass_value = FloatPow::cubed(radius) * 10.;
55
56 let position = Vec3::new(
57 rng.random_range(-1.0..1.0),
58 rng.random_range(-1.0..1.0),
59 rng.random_range(-1.0..1.0),
60 )
61 .normalize()
62 * ops::cbrt(rng.random_range(0.2f32..1.0))
63 * 15.;
64
65 commands.spawn((
66 BodyBundle {
67 mesh: Mesh3d(mesh.clone()),
68 material: MeshMaterial3d(materials.add(Color::srgb(
69 rng.random_range(color_range.clone()),
70 rng.random_range(color_range.clone()),
71 rng.random_range(color_range.clone()),
72 ))),
73 mass: Mass(mass_value),
74 acceleration: Acceleration(Vec3::ZERO),
75 last_pos: LastPos(
76 position
77 - Vec3::new(
78 rng.random_range(vel_range.clone()),
79 rng.random_range(vel_range.clone()),
80 rng.random_range(vel_range.clone()),
81 ) * time.timestep().as_secs_f32(),
82 ),
83 },
84 Transform {
85 translation: position,
86 scale: Vec3::splat(radius),
87 ..default()
88 },
89 ));
90 }
91
92 // add bigger "star" body in the center
93 let star_radius = 1.;
94 commands
95 .spawn((
96 BodyBundle {
97 mesh: Mesh3d(meshes.add(Sphere::new(1.0).mesh().ico(5).unwrap())),
98 material: MeshMaterial3d(materials.add(StandardMaterial {
99 base_color: ORANGE_RED.into(),
100 emissive: LinearRgba::from(ORANGE_RED) * 2.,
101 ..default()
102 })),
103
104 mass: Mass(500.0),
105 ..default()
106 },
107 Transform::from_scale(Vec3::splat(star_radius)),
108 Star,
109 ))
110 .with_child(PointLight {
111 color: Color::WHITE,
112 range: 100.0,
113 radius: star_radius,
114 ..default()
115 });
116 commands.spawn((
117 Camera3d::default(),
118 Transform::from_xyz(0.0, 10.5, -30.0).looking_at(Vec3::ZERO, Vec3::Y),
119 ));
120 }
121
122 fn interact_bodies(mut query: Query<(&Mass, &GlobalTransform, &mut Acceleration)>) {
123 let mut iter = query.iter_combinations_mut();
124 while let Some([(Mass(m1), transform1, mut acc1), (Mass(m2), transform2, mut acc2)]) =
125 iter.fetch_next()
126 {
127 let delta = transform2.translation() - transform1.translation();
128 let distance_sq: f32 = delta.length_squared();
129
130 let f = GRAVITY_CONSTANT / distance_sq;
131 let force_unit_mass = delta * f;
132 acc1.0 += force_unit_mass * *m2;
133 acc2.0 -= force_unit_mass * *m1;
134 }
135 }
136
137 fn integrate(time: Res<Time>, mut query: Query<(&mut Acceleration, &mut Transform, &mut LastPos)>) {
138 let dt_sq = time.delta_secs() * time.delta_secs();
139 for (mut acceleration, mut transform, mut last_pos) in &mut query {
140 // verlet integration
141 // x(t+dt) = 2x(t) - x(t-dt) + a(t)dt^2 + O(dt^4)
142
143 let new_pos = transform.translation * 2.0 - last_pos.0 + acceleration.0 * dt_sq;
144 acceleration.0 = Vec3::ZERO;
145 last_pos.0 = transform.translation;
146 transform.translation = new_pos;
147 }
148 }
149
150 fn look_at_star(
151 mut camera: Single<&mut Transform, (With<Camera>, Without<Star>)>,
152 star: Single<&Transform, With<Star>>,
153 ) {
154 let new_rotation = camera
155 .looking_at(star.translation, Vec3::Y)
156 .rotation
157 .lerp(camera.rotation, 0.1);
158 camera.rotation = new_rotation;
159 }