1use crate::orbit::physics::{NBodyState, OrbitBody};
16use crate::orbit::units::{OrbitMass, Position3D, Velocity3D, AU, EARTH_MASS, G, SOLAR_MASS};
17use serde::{Deserialize, Serialize};
18
19#[derive(Debug, Clone, Serialize, Deserialize)]
21pub enum ScenarioType {
22 Kepler(KeplerConfig),
24 NBody(NBodyConfig),
26 Hohmann(HohmannConfig),
28 Lagrange(LagrangeConfig),
30}
31
32#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct KeplerConfig {
35 pub central_mass: f64,
37 pub orbiter_mass: f64,
39 pub semi_major_axis: f64,
41 pub eccentricity: f64,
43 pub initial_anomaly: f64,
45}
46
47impl Default for KeplerConfig {
48 fn default() -> Self {
49 Self {
50 central_mass: SOLAR_MASS,
51 orbiter_mass: EARTH_MASS,
52 semi_major_axis: AU,
53 eccentricity: 0.0167, initial_anomaly: 0.0,
55 }
56 }
57}
58
59impl KeplerConfig {
60 #[must_use]
62 pub fn earth_sun() -> Self {
63 Self::default()
64 }
65
66 #[must_use]
68 pub fn circular(central_mass: f64, orbiter_mass: f64, radius: f64) -> Self {
69 Self {
70 central_mass,
71 orbiter_mass,
72 semi_major_axis: radius,
73 eccentricity: 0.0,
74 initial_anomaly: 0.0,
75 }
76 }
77
78 #[must_use]
80 pub fn period(&self) -> f64 {
81 2.0 * std::f64::consts::PI * (self.semi_major_axis.powi(3) / (G * self.central_mass)).sqrt()
82 }
83
84 #[must_use]
86 pub fn circular_velocity(&self) -> f64 {
87 (G * self.central_mass / self.semi_major_axis).sqrt()
88 }
89
90 #[must_use]
92 #[allow(clippy::many_single_char_names)] pub fn build(&self, softening: f64) -> NBodyState {
94 let mu = G * self.central_mass;
95
96 let semi_latus = self.semi_major_axis * (1.0 - self.eccentricity * self.eccentricity);
99 let radius = semi_latus / (1.0 + self.eccentricity * self.initial_anomaly.cos());
100
101 let pos_x = radius * self.initial_anomaly.cos();
103 let pos_y = radius * self.initial_anomaly.sin();
104
105 let ang_momentum = (mu * semi_latus).sqrt();
107 let vel_x = -mu / ang_momentum * self.initial_anomaly.sin();
108 let vel_y = mu / ang_momentum * (self.eccentricity + self.initial_anomaly.cos());
109
110 let bodies = vec![
111 OrbitBody::new(
112 OrbitMass::from_kg(self.central_mass),
113 Position3D::zero(),
114 Velocity3D::zero(),
115 ),
116 OrbitBody::new(
117 OrbitMass::from_kg(self.orbiter_mass),
118 Position3D::from_meters(pos_x, pos_y, 0.0),
119 Velocity3D::from_mps(vel_x, vel_y, 0.0),
120 ),
121 ];
122
123 NBodyState::new(bodies, softening)
124 }
125}
126
127#[derive(Debug, Clone, Serialize, Deserialize)]
129pub struct NBodyConfig {
130 pub bodies: Vec<BodyConfig>,
132}
133
134#[derive(Debug, Clone, Serialize, Deserialize)]
136pub struct BodyConfig {
137 pub name: String,
138 pub mass: f64,
139 pub position: (f64, f64, f64),
140 pub velocity: (f64, f64, f64),
141}
142
143impl Default for NBodyConfig {
144 fn default() -> Self {
145 Self::inner_solar_system()
146 }
147}
148
149impl NBodyConfig {
150 #[must_use]
152 pub fn inner_solar_system() -> Self {
153 let sun = BodyConfig {
155 name: "Sun".to_string(),
156 mass: SOLAR_MASS,
157 position: (0.0, 0.0, 0.0),
158 velocity: (0.0, 0.0, 0.0),
159 };
160
161 let mercury = {
162 let r = 0.387 * AU;
163 let v = (G * SOLAR_MASS / r).sqrt();
164 BodyConfig {
165 name: "Mercury".to_string(),
166 mass: 3.301e23,
167 position: (r, 0.0, 0.0),
168 velocity: (0.0, v, 0.0),
169 }
170 };
171
172 let venus = {
173 let r = 0.723 * AU;
174 let v = (G * SOLAR_MASS / r).sqrt();
175 BodyConfig {
176 name: "Venus".to_string(),
177 mass: 4.867e24,
178 position: (0.0, r, 0.0),
179 velocity: (-v, 0.0, 0.0),
180 }
181 };
182
183 let earth = {
184 let r = AU;
185 let v = (G * SOLAR_MASS / r).sqrt();
186 BodyConfig {
187 name: "Earth".to_string(),
188 mass: EARTH_MASS,
189 position: (-r, 0.0, 0.0),
190 velocity: (0.0, -v, 0.0),
191 }
192 };
193
194 let mars = {
195 let r = 1.524 * AU;
196 let v = (G * SOLAR_MASS / r).sqrt();
197 BodyConfig {
198 name: "Mars".to_string(),
199 mass: 6.417e23,
200 position: (0.0, -r, 0.0),
201 velocity: (v, 0.0, 0.0),
202 }
203 };
204
205 Self {
206 bodies: vec![sun, mercury, venus, earth, mars],
207 }
208 }
209
210 #[must_use]
212 pub fn build(&self, softening: f64) -> NBodyState {
213 let bodies = self
214 .bodies
215 .iter()
216 .map(|b| {
217 OrbitBody::new(
218 OrbitMass::from_kg(b.mass),
219 Position3D::from_meters(b.position.0, b.position.1, b.position.2),
220 Velocity3D::from_mps(b.velocity.0, b.velocity.1, b.velocity.2),
221 )
222 })
223 .collect();
224
225 NBodyState::new(bodies, softening)
226 }
227}
228
229#[derive(Debug, Clone, Serialize, Deserialize)]
231pub struct HohmannConfig {
232 pub central_mass: f64,
234 pub spacecraft_mass: f64,
236 pub r1: f64,
238 pub r2: f64,
240 pub phase: usize,
242}
243
244impl Default for HohmannConfig {
245 fn default() -> Self {
246 Self::earth_to_mars()
247 }
248}
249
250impl HohmannConfig {
251 #[must_use]
253 pub fn earth_to_mars() -> Self {
254 Self {
255 central_mass: SOLAR_MASS,
256 spacecraft_mass: 1000.0,
257 r1: AU,
258 r2: 1.524 * AU,
259 phase: 0,
260 }
261 }
262
263 #[must_use]
265 pub fn leo_to_geo() -> Self {
266 let _earth_mu = 3.986e14; let r_earth = 6.378e6; Self {
269 central_mass: EARTH_MASS,
270 spacecraft_mass: 1000.0,
271 r1: r_earth + 400_000.0, r2: r_earth + 35_786_000.0, phase: 0,
274 }
275 }
276
277 #[must_use]
279 pub fn delta_v1(&self) -> f64 {
280 let mu = G * self.central_mass;
281 let v_circular = (mu / self.r1).sqrt();
282 let v_transfer = (2.0 * mu * self.r2 / (self.r1 * (self.r1 + self.r2))).sqrt();
283 v_transfer - v_circular
284 }
285
286 #[must_use]
288 pub fn delta_v2(&self) -> f64 {
289 let mu = G * self.central_mass;
290 let v_transfer = (2.0 * mu * self.r1 / (self.r2 * (self.r1 + self.r2))).sqrt();
291 let v_circular = (mu / self.r2).sqrt();
292 v_circular - v_transfer
293 }
294
295 #[must_use]
297 pub fn total_delta_v(&self) -> f64 {
298 self.delta_v1().abs() + self.delta_v2().abs()
299 }
300
301 #[must_use]
303 pub fn transfer_time(&self) -> f64 {
304 let mu = G * self.central_mass;
305 let a = (self.r1 + self.r2) / 2.0;
306 std::f64::consts::PI * (a.powi(3) / mu).sqrt()
307 }
308
309 #[must_use]
311 pub fn build_initial(&self, softening: f64) -> NBodyState {
312 let mu = G * self.central_mass;
313 let v = (mu / self.r1).sqrt();
314
315 let bodies = vec![
316 OrbitBody::new(
317 OrbitMass::from_kg(self.central_mass),
318 Position3D::zero(),
319 Velocity3D::zero(),
320 ),
321 OrbitBody::new(
322 OrbitMass::from_kg(self.spacecraft_mass),
323 Position3D::from_meters(self.r1, 0.0, 0.0),
324 Velocity3D::from_mps(0.0, v, 0.0),
325 ),
326 ];
327
328 NBodyState::new(bodies, softening)
329 }
330
331 #[must_use]
333 pub fn build_transfer(&self, softening: f64) -> NBodyState {
334 let mu = G * self.central_mass;
335 let v = (mu / self.r1).sqrt() + self.delta_v1();
336
337 let bodies = vec![
338 OrbitBody::new(
339 OrbitMass::from_kg(self.central_mass),
340 Position3D::zero(),
341 Velocity3D::zero(),
342 ),
343 OrbitBody::new(
344 OrbitMass::from_kg(self.spacecraft_mass),
345 Position3D::from_meters(self.r1, 0.0, 0.0),
346 Velocity3D::from_mps(0.0, v, 0.0),
347 ),
348 ];
349
350 NBodyState::new(bodies, softening)
351 }
352}
353
354#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
356pub enum LagrangePoint {
357 L1,
358 L2,
359 L3,
360 L4,
361 L5,
362}
363
364#[derive(Debug, Clone, Serialize, Deserialize)]
366pub struct LagrangeConfig {
367 pub primary_mass: f64,
369 pub secondary_mass: f64,
371 pub separation: f64,
373 pub point: LagrangePoint,
375 pub perturbation: (f64, f64, f64),
377}
378
379impl Default for LagrangeConfig {
380 fn default() -> Self {
381 Self::sun_earth_l2()
382 }
383}
384
385impl LagrangeConfig {
386 #[must_use]
388 pub fn sun_earth_l2() -> Self {
389 Self {
390 primary_mass: SOLAR_MASS,
391 secondary_mass: EARTH_MASS,
392 separation: AU,
393 point: LagrangePoint::L2,
394 perturbation: (0.0, 0.0, 0.0),
395 }
396 }
397
398 #[must_use]
400 pub fn sun_earth_l4() -> Self {
401 Self {
402 primary_mass: SOLAR_MASS,
403 secondary_mass: EARTH_MASS,
404 separation: AU,
405 point: LagrangePoint::L4,
406 perturbation: (0.0, 0.0, 0.0),
407 }
408 }
409
410 #[must_use]
412 pub fn mass_ratio(&self) -> f64 {
413 self.secondary_mass / (self.primary_mass + self.secondary_mass)
414 }
415
416 #[must_use]
418 pub fn lagrange_position(&self) -> (f64, f64, f64) {
419 let mu = self.mass_ratio();
420 let r = self.separation;
421
422 match self.point {
423 LagrangePoint::L1 => {
424 let x = r * (1.0 - (mu / 3.0).cbrt());
426 (x, 0.0, 0.0)
427 }
428 LagrangePoint::L2 => {
429 let x = r * (1.0 + (mu / 3.0).cbrt());
431 (x, 0.0, 0.0)
432 }
433 LagrangePoint::L3 => {
434 let x = -r * (1.0 + 5.0 * mu / 12.0);
436 (x, 0.0, 0.0)
437 }
438 LagrangePoint::L4 => {
439 let x = r * 0.5;
441 let y = r * (3.0_f64.sqrt() / 2.0);
442 (x, y, 0.0)
443 }
444 LagrangePoint::L5 => {
445 let x = r * 0.5;
447 let y = -r * (3.0_f64.sqrt() / 2.0);
448 (x, y, 0.0)
449 }
450 }
451 }
452
453 #[must_use]
455 pub fn build(&self, softening: f64) -> NBodyState {
456 let mu = self.mass_ratio();
457 let total_mass = self.primary_mass + self.secondary_mass;
458 let r = self.separation;
459
460 let x1 = -mu * r; let x2 = (1.0 - mu) * r; let omega = (G * total_mass / r.powi(3)).sqrt();
466
467 let v2 = omega * x2.abs();
469
470 let (lx, ly, lz) = self.lagrange_position();
472 let lx = lx + self.perturbation.0;
473 let ly = ly + self.perturbation.1;
474 let lz = lz + self.perturbation.2;
475
476 let test_r = (lx * lx + ly * ly).sqrt();
478 let test_v = omega * test_r;
479 let test_vx = -test_v * ly / test_r.max(1e-10);
480 let test_vy = test_v * lx / test_r.max(1e-10);
481
482 let bodies = vec![
483 OrbitBody::new(
484 OrbitMass::from_kg(self.primary_mass),
485 Position3D::from_meters(x1, 0.0, 0.0),
486 Velocity3D::zero(),
487 ),
488 OrbitBody::new(
489 OrbitMass::from_kg(self.secondary_mass),
490 Position3D::from_meters(x2, 0.0, 0.0),
491 Velocity3D::from_mps(0.0, v2, 0.0),
492 ),
493 OrbitBody::new(
494 OrbitMass::from_kg(1.0), Position3D::from_meters(lx, ly, lz),
496 Velocity3D::from_mps(test_vx, test_vy, 0.0),
497 ),
498 ];
499
500 NBodyState::new(bodies, softening)
501 }
502}
503
504#[cfg(test)]
505mod tests {
506 use super::*;
507
508 const EPSILON: f64 = 1e-6;
509
510 #[test]
511 fn test_kepler_config_default() {
512 let config = KeplerConfig::default();
513 assert!((config.central_mass - SOLAR_MASS).abs() < 1e20);
514 assert!((config.orbiter_mass - EARTH_MASS).abs() < 1e20);
515 }
516
517 #[test]
518 fn test_kepler_period() {
519 let config = KeplerConfig::earth_sun();
520 let period = config.period();
521 let expected = 365.25 * 86400.0; let relative_error = (period - expected).abs() / expected;
523 assert!(relative_error < 0.001, "Period error: {relative_error}");
524 }
525
526 #[test]
527 fn test_kepler_circular_velocity() {
528 let config = KeplerConfig::earth_sun();
529 let v = config.circular_velocity();
530 let expected = 29780.0; let relative_error = (v - expected).abs() / expected;
532 assert!(relative_error < 0.001, "Velocity error: {relative_error}");
533 }
534
535 #[test]
536 fn test_kepler_build() {
537 let config = KeplerConfig::earth_sun();
538 let state = config.build(0.0);
539 assert_eq!(state.num_bodies(), 2);
540 }
541
542 #[test]
543 fn test_nbody_inner_solar_system() {
544 let config = NBodyConfig::inner_solar_system();
545 assert_eq!(config.bodies.len(), 5);
546 assert_eq!(config.bodies[0].name, "Sun");
547 }
548
549 #[test]
550 fn test_nbody_build() {
551 let config = NBodyConfig::inner_solar_system();
552 let state = config.build(1e6);
553 assert_eq!(state.num_bodies(), 5);
554 }
555
556 #[test]
557 fn test_hohmann_delta_v1() {
558 let config = HohmannConfig::earth_to_mars();
559 let dv1 = config.delta_v1();
560 assert!(dv1 > 0.0); assert!(dv1 < 5000.0); }
563
564 #[test]
565 fn test_hohmann_delta_v2() {
566 let config = HohmannConfig::earth_to_mars();
567 let dv2 = config.delta_v2();
568 assert!(dv2 > 0.0); assert!(dv2 < 5000.0);
570 }
571
572 #[test]
573 fn test_hohmann_total_delta_v() {
574 let config = HohmannConfig::earth_to_mars();
575 let total = config.total_delta_v();
576 let expected = 5600.0; let relative_error = (total - expected).abs() / expected;
578 assert!(
579 relative_error < 0.1,
580 "Total delta-v error: {relative_error}"
581 );
582 }
583
584 #[test]
585 fn test_hohmann_transfer_time() {
586 let config = HohmannConfig::earth_to_mars();
587 let time = config.transfer_time();
588 let expected = 259.0 * 86400.0; let relative_error = (time - expected).abs() / expected;
590 assert!(
591 relative_error < 0.1,
592 "Transfer time error: {relative_error}"
593 );
594 }
595
596 #[test]
597 fn test_hohmann_build_initial() {
598 let config = HohmannConfig::earth_to_mars();
599 let state = config.build_initial(0.0);
600 assert_eq!(state.num_bodies(), 2);
601 }
602
603 #[test]
604 fn test_lagrange_mass_ratio() {
605 let config = LagrangeConfig::sun_earth_l2();
606 let mu = config.mass_ratio();
607 assert!(mu > 0.0);
608 assert!(mu < 1e-5); }
610
611 #[test]
612 fn test_lagrange_l1_position() {
613 let config = LagrangeConfig {
614 point: LagrangePoint::L1,
615 ..LagrangeConfig::default()
616 };
617 let (x, y, z) = config.lagrange_position();
618 assert!(x > 0.0); assert!(x < AU);
620 assert!(y.abs() < EPSILON);
621 assert!(z.abs() < EPSILON);
622 }
623
624 #[test]
625 fn test_lagrange_l2_position() {
626 let config = LagrangeConfig::sun_earth_l2();
627 let (x, y, z) = config.lagrange_position();
628 assert!(x > AU); assert!(y.abs() < EPSILON);
630 assert!(z.abs() < EPSILON);
631 }
632
633 #[test]
634 fn test_lagrange_l4_position() {
635 let config = LagrangeConfig::sun_earth_l4();
636 let (x, y, z) = config.lagrange_position();
637 assert!((x - 0.5 * AU).abs() / AU < 0.01);
639 assert!((y - 0.866 * AU).abs() / AU < 0.01);
640 assert!(z.abs() < EPSILON);
641 }
642
643 #[test]
644 fn test_lagrange_build() {
645 let config = LagrangeConfig::sun_earth_l2();
646 let state = config.build(1e6);
647 assert_eq!(state.num_bodies(), 3); }
649
650 #[test]
651 fn test_scenario_type_kepler() {
652 let scenario = ScenarioType::Kepler(KeplerConfig::default());
653 match scenario {
654 ScenarioType::Kepler(config) => {
655 assert!((config.eccentricity - 0.0167).abs() < 0.001);
656 }
657 _ => panic!("Expected Kepler scenario"),
658 }
659 }
660
661 #[test]
662 fn test_kepler_circular() {
663 let config = KeplerConfig::circular(SOLAR_MASS, EARTH_MASS, AU);
664 assert!((config.eccentricity).abs() < EPSILON);
665 assert!((config.central_mass - SOLAR_MASS).abs() < 1e20);
666 }
667
668 #[test]
669 fn test_nbody_default() {
670 let config = NBodyConfig::default();
671 assert_eq!(config.bodies.len(), 5);
672 }
673
674 #[test]
675 fn test_hohmann_default() {
676 let config = HohmannConfig::default();
677 assert!((config.r1 - AU).abs() < 1e6);
678 }
679
680 #[test]
681 fn test_hohmann_leo_to_geo() {
682 let config = HohmannConfig::leo_to_geo();
683 assert!(config.r1 > 6e6); assert!(config.r2 > 4e7); assert!(config.central_mass > 5e24); }
687
688 #[test]
689 fn test_hohmann_build_transfer() {
690 let config = HohmannConfig::earth_to_mars();
691 let state = config.build_transfer(0.0);
692 assert_eq!(state.num_bodies(), 2);
693
694 let (x, _, _) = state.bodies[1].position.as_meters();
696 assert!((x - config.r1).abs() < 1e6);
697 }
698
699 #[test]
700 fn test_lagrange_default() {
701 let config = LagrangeConfig::default();
702 assert_eq!(config.point, LagrangePoint::L2);
703 }
704
705 #[test]
706 fn test_lagrange_l3_position() {
707 let config = LagrangeConfig {
708 point: LagrangePoint::L3,
709 ..LagrangeConfig::default()
710 };
711 let (x, y, z) = config.lagrange_position();
712 assert!(x < 0.0); assert!(x < -AU * 0.9);
714 assert!(y.abs() < EPSILON);
715 assert!(z.abs() < EPSILON);
716 }
717
718 #[test]
719 fn test_lagrange_l5_position() {
720 let config = LagrangeConfig {
721 point: LagrangePoint::L5,
722 ..LagrangeConfig::default()
723 };
724 let (x, y, z) = config.lagrange_position();
725 assert!((x - 0.5 * AU).abs() / AU < 0.01);
727 assert!((y + 0.866 * AU).abs() / AU < 0.01); assert!(z.abs() < EPSILON);
729 }
730
731 #[test]
732 fn test_lagrange_with_perturbation() {
733 let config = LagrangeConfig {
734 perturbation: (1000.0, 2000.0, 3000.0),
735 ..LagrangeConfig::sun_earth_l2()
736 };
737 let state = config.build(1e6);
738 assert_eq!(state.num_bodies(), 3);
739
740 let (lx, _, _) = state.bodies[2].position.as_meters();
742 let nominal_l2 = config.lagrange_position().0;
743 assert!((lx - nominal_l2 - 1000.0).abs() < 1.0);
744 }
745
746 #[test]
747 fn test_scenario_type_nbody() {
748 let scenario = ScenarioType::NBody(NBodyConfig::default());
749 match scenario {
750 ScenarioType::NBody(config) => {
751 assert_eq!(config.bodies.len(), 5);
752 }
753 _ => panic!("Expected NBody scenario"),
754 }
755 }
756
757 #[test]
758 fn test_scenario_type_hohmann() {
759 let scenario = ScenarioType::Hohmann(HohmannConfig::default());
760 match scenario {
761 ScenarioType::Hohmann(config) => {
762 assert!((config.r1 - AU).abs() < 1e6);
763 }
764 _ => panic!("Expected Hohmann scenario"),
765 }
766 }
767
768 #[test]
769 fn test_scenario_type_lagrange() {
770 let scenario = ScenarioType::Lagrange(LagrangeConfig::default());
771 match scenario {
772 ScenarioType::Lagrange(config) => {
773 assert_eq!(config.point, LagrangePoint::L2);
774 }
775 _ => panic!("Expected Lagrange scenario"),
776 }
777 }
778
779 #[test]
780 fn test_kepler_build_with_nonzero_anomaly() {
781 let config = KeplerConfig {
782 initial_anomaly: std::f64::consts::PI / 2.0, ..KeplerConfig::default()
784 };
785 let state = config.build(0.0);
786 assert_eq!(state.num_bodies(), 2);
787
788 let (_, y, _) = state.bodies[1].position.as_meters();
790 assert!(y.abs() > AU * 0.5);
791 }
792
793 #[test]
794 fn test_body_config_fields() {
795 let body = BodyConfig {
796 name: "TestBody".to_string(),
797 mass: 1e24,
798 position: (1e11, 2e11, 3e11),
799 velocity: (1000.0, 2000.0, 3000.0),
800 };
801 assert_eq!(body.name, "TestBody");
802 assert!((body.mass - 1e24).abs() < 1e18);
803 assert!((body.position.0 - 1e11).abs() < 1e6);
804 assert!((body.velocity.0 - 1000.0).abs() < 0.1);
805 }
806}