1extern crate nalgebra;
7
8use crate::atom::*;
9use crate::constant;
10use crate::initiate::NewlyCreated;
11use specs::prelude::*;
12
13pub struct Step {
15 pub n: u64,
16}
17
18pub struct Timestep {
26 pub delta: f64,
28}
29
30pub struct EulerIntegrationSystem;
38
39impl<'a> System<'a> for EulerIntegrationSystem {
40 type SystemData = (
41 WriteStorage<'a, Position>,
42 WriteStorage<'a, Velocity>,
43 ReadExpect<'a, Timestep>,
44 WriteExpect<'a, Step>,
45 ReadStorage<'a, Force>,
46 ReadStorage<'a, Mass>,
47 );
48
49 fn run(&mut self, (mut pos, mut vel, t, mut step, force, mass): Self::SystemData) {
50 use rayon::prelude::*;
51
52 step.n += 1;
53 (&mut vel, &mut pos, &force, &mass).par_join().for_each(
54 |(vel, pos, force, mass)| {
55 euler_update(vel, pos, force, mass, t.delta);
56 },
57 );
58 }
59}
60
61pub const INTEGRATE_POSITION_SYSTEM_NAME: &str = "integrate_position";
62
63pub struct VelocityVerletIntegratePositionSystem;
70impl<'a> System<'a> for VelocityVerletIntegratePositionSystem {
71 type SystemData = (
72 WriteStorage<'a, Position>,
73 ReadStorage<'a, Velocity>,
74 ReadExpect<'a, Timestep>,
75 WriteExpect<'a, Step>,
76 ReadStorage<'a, Force>,
77 WriteStorage<'a, OldForce>,
78 ReadStorage<'a, Mass>,
79 );
80
81 fn run(&mut self, (mut pos, vel, t, mut step, force, mut old_force, mass): Self::SystemData) {
82 use rayon::prelude::*;
83
84 step.n += 1;
85 let dt = t.delta;
86
87 (&mut pos, &vel, &mut old_force, &force, &mass)
88 .par_join()
89 .for_each(|(mut pos, vel, mut old_force, force, mass)| {
90 pos.pos = pos.pos
91 + vel.vel * dt
92 + force.force / (constant::AMU * mass.value) / 2.0 * dt * dt;
93 old_force.0 = *force;
94 });
95 }
96}
97
98pub const INTEGRATE_VELOCITY_SYSTEM_NAME: &str = "integrate_velocity";
99
100pub struct VelocityVerletIntegrateVelocitySystem;
106impl<'a> System<'a> for VelocityVerletIntegrateVelocitySystem {
107 type SystemData = (
108 WriteStorage<'a, Velocity>,
109 ReadExpect<'a, Timestep>,
110 ReadStorage<'a, Force>,
111 ReadStorage<'a, OldForce>,
112 ReadStorage<'a, Mass>,
113 );
114
115 fn run(&mut self, (mut vel, t, force, old_force, mass): Self::SystemData) {
116 use rayon::prelude::*;
117
118 let dt = t.delta;
119
120 (&mut vel, &force, &old_force, &mass).par_join().for_each(
121 |(vel, force, old_force, mass)| {
122 vel.vel += (force.force + old_force.0.force) / (constant::AMU * mass.value) / 2.0 * dt;
123 },
124 );
125 }
126}
127
128pub struct AddOldForceToNewAtomsSystem;
130impl<'a> System<'a> for AddOldForceToNewAtomsSystem {
131 type SystemData = (
132 Entities<'a>,
133 ReadStorage<'a, NewlyCreated>,
134 ReadStorage<'a, OldForce>,
135 Read<'a, LazyUpdate>,
136 );
137 fn run(&mut self, (ent, newly_created, old_force, updater): Self::SystemData) {
138 for (ent, _, _) in (&ent, &newly_created, !&old_force).join() {
139 updater.insert(ent, OldForce::default());
140 }
141 }
142}
143
144#[derive(Default)]
146pub struct OldForce(Force);
147impl Component for OldForce {
148 type Storage = VecStorage<OldForce>;
149}
150
151fn euler_update(vel: &mut Velocity, pos: &mut Position, force: &Force, mass: &Mass, dt: f64) {
153 pos.pos += vel.vel * dt;
154 vel.vel += force.force * dt / (constant::AMU * mass.value);
155}
156
157pub mod tests {
158 #[allow(unused_imports)]
159 use super::*;
160 extern crate specs;
161 #[allow(unused_imports)]
162 use specs::{Builder, DispatcherBuilder, World};
163
164 extern crate nalgebra;
165 #[allow(unused_imports)]
166 use nalgebra::Vector3;
167
168 #[test]
169 fn test_euler() {
170 let mut pos = Position {
171 pos: Vector3::new(1., 1., 1.),
172 };
173 let mut vel = Velocity {
174 vel: Vector3::new(0., 1., 0.),
175 };
176 let time = 1.;
177 let mass = Mass {
178 value: 1. / constant::AMU,
179 };
180 let force = Force {
181 force: Vector3::new(1., 1., 1.),
182 };
183 euler_update(&mut vel, &mut pos, &force, &mass, time);
184 assert_eq!(vel.vel, Vector3::new(1., 2., 1.));
185 assert_eq!(pos.pos, Vector3::new(1., 2., 1.));
186 }
187
188 #[test]
190 fn test_euler_integration() {
191 let mut world = World::new();
192
193 let mut dispatcher = DispatcherBuilder::new()
194 .with(EulerIntegrationSystem, "integrator", &[])
195 .build();
196 dispatcher.setup(&mut world);
197
198 let force = Vector3::new(1.0, 0.0, 0.0);
200 let mass = 1.0;
201 let atom = world
202 .create_entity()
203 .with(Position {
204 pos: Vector3::new(0.0, 0.0, 0.0),
205 })
206 .with(Velocity {
207 vel: Vector3::new(0.0, 0.0, 0.0),
208 })
209 .with(Force { force })
210 .with(Mass {
211 value: mass / constant::AMU,
212 })
213 .build();
214
215 let dt = 1.0e-3;
216 world.insert(Timestep { delta: dt });
217 world.insert(Step { n: 0 });
218
219 let n_steps = 1_000;
221 for _i in 0..n_steps {
222 dispatcher.dispatch(&world);
223 world.maintain();
224 }
225
226 let a = force / mass;
227 let expected_v = a * (n_steps as f64 * dt);
228
229 assert_approx_eq::assert_approx_eq!(
230 expected_v.norm(),
231 world
232 .read_storage::<Velocity>()
233 .get(atom)
234 .expect("atom not found.")
235 .vel
236 .norm(),
237 expected_v.norm() * 0.01
238 );
239
240 let expected_x = a * (n_steps as f64 * dt).powi(2) / 2.0;
241 assert_approx_eq::assert_approx_eq!(
242 expected_x.norm(),
243 world
244 .read_storage::<Position>()
245 .get(atom)
246 .expect("atom not found.")
247 .pos
248 .norm(),
249 expected_x.norm() * 0.01
250 );
251 }
252
253 #[test]
254 fn test_add_old_force_system() {
255 let mut test_world = World::new();
256
257 let mut dispatcher = DispatcherBuilder::new()
258 .with(AddOldForceToNewAtomsSystem, "", &[])
259 .build();
260 dispatcher.setup(&mut test_world);
261 test_world.register::<OldForce>();
262
263 let test_entity = test_world.create_entity().with(NewlyCreated {}).build();
264
265 dispatcher.dispatch(&test_world);
266 test_world.maintain();
267
268 let old_forces = test_world.read_storage::<OldForce>();
269 assert!(
270 old_forces.contains(test_entity),
271 "OldForce component not added to test entity."
272 );
273 }
274
275 #[test]
276 fn test_velocity_verlet_integration() {
277 let mut world = World::new();
278
279 let mut dispatcher = DispatcherBuilder::new()
280 .with(
281 VelocityVerletIntegratePositionSystem,
282 "integrate_position",
283 &[],
284 )
285 .with(
286 VelocityVerletIntegrateVelocitySystem,
287 "integrate_velocity",
288 &["integrate_position"],
289 )
290 .build();
291 dispatcher.setup(&mut world);
292
293 let force = Vector3::new(1.0, 0.0, 0.0);
295 let mass = 1.0;
296 let atom = world
297 .create_entity()
298 .with(Position {
299 pos: Vector3::new(0.0, 0.0, 0.0),
300 })
301 .with(Velocity {
302 vel: Vector3::new(0.0, 0.0, 0.0),
303 })
304 .with(Force { force })
305 .with(OldForce {
306 0: Force { force },
307 })
308 .with(Mass {
309 value: mass / constant::AMU,
310 })
311 .build();
312
313 let dt = 1.0e-3;
314 world.insert(Timestep { delta: dt });
315 world.insert(Step { n: 0 });
316
317 let n_steps = 1_000;
319 for _i in 0..n_steps {
320 dispatcher.dispatch(&world);
321 world.maintain();
322 }
323
324 let a = force / mass;
325 let expected_v = a * (n_steps as f64 * dt);
326
327 assert_approx_eq::assert_approx_eq!(
328 expected_v.norm(),
329 world
330 .read_storage::<Velocity>()
331 .get(atom)
332 .expect("atom not found.")
333 .vel
334 .norm(),
335 expected_v.norm() * 0.01
336 );
337
338 let expected_x = a * (n_steps as f64 * dt).powi(2) / 2.0;
339 assert_approx_eq::assert_approx_eq!(
340 expected_x.norm(),
341 world
342 .read_storage::<Position>()
343 .get(atom)
344 .expect("atom not found.")
345 .pos
346 .norm(),
347 expected_x.norm() * 0.01
348 );
349 }
350}