1pub mod pontryagin;
30pub mod zoh;
32
33use crate::error::ensure_finite;
34use crate::integration::{
35 DifferentiableDynamicsModel, Dop853, DynamicsModel, InitialValueProblem, IntegratorOptions,
36 Propagation, SensitivityProblem, SensitivityPropagation,
37};
38use crate::{CartesianState, Matrix6, PykepError, Result};
39
40const POSITION_DIMENSION: usize = 3;
41
42#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
44pub struct KeplerDynamics;
45
46impl KeplerDynamics {
47 pub fn evaluate(&self, state: &CartesianState, mu: f64) -> Result<CartesianState> {
55 let mut derivative = [0.0; 6];
56 self.rhs(0.0, state, &[mu], &mut derivative)?;
57 Ok(derivative)
58 }
59
60 pub fn propagate(
66 &self,
67 initial_time: f64,
68 initial_state: CartesianState,
69 final_time: f64,
70 mu: f64,
71 options: IntegratorOptions,
72 ) -> Result<Propagation<6>> {
73 Dop853.propagate(
74 self,
75 InitialValueProblem::new(initial_time, initial_state, final_time, [mu]),
76 options,
77 )
78 }
79
80 pub fn propagate_with_stm(
86 &self,
87 initial_time: f64,
88 initial_state: CartesianState,
89 final_time: f64,
90 mu: f64,
91 options: IntegratorOptions,
92 ) -> Result<SensitivityPropagation<6, 6>> {
93 propagate_stm(self, initial_time, initial_state, final_time, [mu], options)
94 }
95}
96
97impl DynamicsModel<6, 1> for KeplerDynamics {
98 const NAME: &'static str = "Kepler dynamics";
99
100 fn validate(&self, time: f64, state: &CartesianState, parameters: &[f64; 1]) -> Result<()> {
101 validate_time_state(time, state)?;
102 validate_positive("mu", parameters[0])?;
103 radius_squared([state[0], state[1], state[2]], "Kepler dynamics radius")?;
104 Ok(())
105 }
106
107 fn rhs(
108 &self,
109 time: f64,
110 state: &CartesianState,
111 parameters: &[f64; 1],
112 derivative: &mut CartesianState,
113 ) -> Result<()> {
114 self.validate(time, state, parameters)?;
115 let position = [state[0], state[1], state[2]];
116 let (acceleration, _) = point_mass_acceleration_and_gradient(
117 position,
118 parameters[0],
119 "Kepler dynamics radius",
120 )?;
121 *derivative = [
122 state[3],
123 state[4],
124 state[5],
125 acceleration[0],
126 acceleration[1],
127 acceleration[2],
128 ];
129 validate_output(Self::NAME, derivative)
130 }
131}
132
133impl DifferentiableDynamicsModel<6, 1> for KeplerDynamics {
134 fn jacobians(
135 &self,
136 time: f64,
137 state: &CartesianState,
138 parameters: &[f64; 1],
139 state_jacobian: &mut Matrix6,
140 parameter_jacobian: &mut [[f64; 1]; 6],
141 ) -> Result<()> {
142 self.validate(time, state, parameters)?;
143 let position = [state[0], state[1], state[2]];
144 let (_, gradient) = point_mass_acceleration_and_gradient(
145 position,
146 parameters[0],
147 "Kepler dynamics radius",
148 )?;
149 *state_jacobian = kinematic_jacobian();
150 set_spatial_gradient(state_jacobian, &gradient);
151 let radius = radius_squared(position, "Kepler dynamics radius")?.sqrt();
152 let inverse_radius_cubed = 1.0 / (radius * radius * radius);
153 *parameter_jacobian = [[0.0]; 6];
154 for row in 0..POSITION_DIMENSION {
155 parameter_jacobian[row + 3][0] = -position[row] * inverse_radius_cubed;
156 }
157 validate_jacobians(Self::NAME, state_jacobian, parameter_jacobian)
158 }
159}
160
161#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
163pub struct Cr3bpDynamics;
164
165impl Cr3bpDynamics {
166 pub fn evaluate(&self, state: &CartesianState, mu: f64) -> Result<CartesianState> {
176 let mut derivative = [0.0; 6];
177 self.rhs(0.0, state, &[mu], &mut derivative)?;
178 Ok(derivative)
179 }
180
181 pub fn effective_potential(&self, state: &CartesianState, mu: f64) -> Result<f64> {
188 self.validate(0.0, state, &[mu])?;
189 let (d1, d2) = primary_displacements(state, mu);
190 let r1 = radius_squared(d1, "CR3BP primary-one distance")?.sqrt();
191 let r2 = radius_squared(d2, "CR3BP primary-two distance")?.sqrt();
192 let potential =
193 0.5 * (state[0] * state[0] + state[1] * state[1]) + (1.0 - mu) / r1 + mu / r2;
194 finite_output(Self::NAME, potential)
195 }
196
197 pub fn jacobi_constant(&self, state: &CartesianState, mu: f64) -> Result<f64> {
203 let potential = self.effective_potential(state, mu)?;
204 let velocity_squared = state[3] * state[3] + state[4] * state[4] + state[5] * state[5];
205 finite_output(Self::NAME, 2.0 * potential - velocity_squared)
206 }
207
208 pub fn propagate(
214 &self,
215 initial_time: f64,
216 initial_state: CartesianState,
217 final_time: f64,
218 mu: f64,
219 options: IntegratorOptions,
220 ) -> Result<Propagation<6>> {
221 Dop853.propagate(
222 self,
223 InitialValueProblem::new(initial_time, initial_state, final_time, [mu]),
224 options,
225 )
226 }
227
228 pub fn propagate_with_stm(
234 &self,
235 initial_time: f64,
236 initial_state: CartesianState,
237 final_time: f64,
238 mu: f64,
239 options: IntegratorOptions,
240 ) -> Result<SensitivityPropagation<6, 6>> {
241 propagate_stm(self, initial_time, initial_state, final_time, [mu], options)
242 }
243}
244
245impl DynamicsModel<6, 1> for Cr3bpDynamics {
246 const NAME: &'static str = "CR3BP dynamics";
247
248 fn validate(&self, time: f64, state: &CartesianState, parameters: &[f64; 1]) -> Result<()> {
249 validate_time_state(time, state)?;
250 validate_mass_fraction(parameters[0])?;
251 let (d1, d2) = primary_displacements(state, parameters[0]);
252 radius_squared(d1, "CR3BP primary-one distance")?;
253 radius_squared(d2, "CR3BP primary-two distance")?;
254 Ok(())
255 }
256
257 fn rhs(
258 &self,
259 time: f64,
260 state: &CartesianState,
261 parameters: &[f64; 1],
262 derivative: &mut CartesianState,
263 ) -> Result<()> {
264 self.validate(time, state, parameters)?;
265 let mu = parameters[0];
266 let (d1, d2) = primary_displacements(state, mu);
267 let (a1, _) =
268 point_mass_acceleration_and_gradient(d1, 1.0 - mu, "CR3BP primary-one distance")?;
269 let (a2, _) = point_mass_acceleration_and_gradient(d2, mu, "CR3BP primary-two distance")?;
270 *derivative = [
271 state[3],
272 state[4],
273 state[5],
274 2.0 * state[4] + state[0] + a1[0] + a2[0],
275 -2.0 * state[3] + state[1] + a1[1] + a2[1],
276 a1[2] + a2[2],
277 ];
278 validate_output(Self::NAME, derivative)
279 }
280}
281
282impl DifferentiableDynamicsModel<6, 1> for Cr3bpDynamics {
283 fn jacobians(
284 &self,
285 time: f64,
286 state: &CartesianState,
287 parameters: &[f64; 1],
288 state_jacobian: &mut Matrix6,
289 parameter_jacobian: &mut [[f64; 1]; 6],
290 ) -> Result<()> {
291 self.validate(time, state, parameters)?;
292 let mu = parameters[0];
293 let (d1, d2) = primary_displacements(state, mu);
294 let (_, gradient1) =
295 point_mass_acceleration_and_gradient(d1, 1.0 - mu, "CR3BP primary-one distance")?;
296 let (_, gradient2) =
297 point_mass_acceleration_and_gradient(d2, mu, "CR3BP primary-two distance")?;
298 *state_jacobian = rotating_kinematic_jacobian();
299 add_spatial_gradient(state_jacobian, &gradient1);
300 add_spatial_gradient(state_jacobian, &gradient2);
301
302 *parameter_jacobian = [[0.0]; 6];
303 let derivative1 =
304 moving_mass_parameter_derivative(d1, 1.0 - mu, -1.0, "CR3BP primary-one distance")?;
305 let derivative2 =
306 moving_mass_parameter_derivative(d2, mu, 1.0, "CR3BP primary-two distance")?;
307 for row in 0..POSITION_DIMENSION {
308 parameter_jacobian[row + 3][0] = derivative1[row] + derivative2[row];
309 }
310 validate_jacobians(Self::NAME, state_jacobian, parameter_jacobian)
311 }
312}
313
314#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
316pub struct BcpDynamics;
317
318impl BcpDynamics {
319 pub fn evaluate(
328 &self,
329 time: f64,
330 state: &CartesianState,
331 parameters: [f64; 4],
332 ) -> Result<CartesianState> {
333 let mut derivative = [0.0; 6];
334 self.rhs(time, state, ¶meters, &mut derivative)?;
335 Ok(derivative)
336 }
337
338 pub fn propagate(
344 &self,
345 initial_time: f64,
346 initial_state: CartesianState,
347 final_time: f64,
348 parameters: [f64; 4],
349 options: IntegratorOptions,
350 ) -> Result<Propagation<6>> {
351 Dop853.propagate(
352 self,
353 InitialValueProblem::new(initial_time, initial_state, final_time, parameters),
354 options,
355 )
356 }
357
358 pub fn propagate_with_stm(
364 &self,
365 initial_time: f64,
366 initial_state: CartesianState,
367 final_time: f64,
368 parameters: [f64; 4],
369 options: IntegratorOptions,
370 ) -> Result<SensitivityPropagation<6, 6>> {
371 propagate_stm(
372 self,
373 initial_time,
374 initial_state,
375 final_time,
376 parameters,
377 options,
378 )
379 }
380}
381
382impl DynamicsModel<6, 4> for BcpDynamics {
383 const NAME: &'static str = "BCP dynamics";
384
385 fn validate(&self, time: f64, state: &CartesianState, parameters: &[f64; 4]) -> Result<()> {
386 validate_time_state(time, state)?;
387 validate_mass_fraction(parameters[0])?;
388 validate_non_negative("mu_sun", parameters[1])?;
389 validate_positive("rho_sun", parameters[2])?;
390 ensure_finite("omega_sun", parameters[3])?;
391 let (d1, d2) = primary_displacements(state, parameters[0]);
392 radius_squared(d1, "BCP primary-one distance")?;
393 radius_squared(d2, "BCP primary-two distance")?;
394 let (_, _, sun_displacement) = sun_geometry(time, state, parameters[2], parameters[3]);
395 radius_squared(sun_displacement, "BCP Sun distance")?;
396 Ok(())
397 }
398
399 fn rhs(
400 &self,
401 time: f64,
402 state: &CartesianState,
403 parameters: &[f64; 4],
404 derivative: &mut CartesianState,
405 ) -> Result<()> {
406 self.validate(time, state, parameters)?;
407 let [mu, mu_sun, rho_sun, omega_sun] = *parameters;
408 let (d1, d2) = primary_displacements(state, mu);
409 let (a1, _) =
410 point_mass_acceleration_and_gradient(d1, 1.0 - mu, "BCP primary-one distance")?;
411 let (a2, _) = point_mass_acceleration_and_gradient(d2, mu, "BCP primary-two distance")?;
412 let (sun_direction, _, sun_displacement) = sun_geometry(time, state, rho_sun, omega_sun);
413 let (sun_acceleration, _) =
414 point_mass_acceleration_and_gradient(sun_displacement, mu_sun, "BCP Sun distance")?;
415 let indirect_scale = -mu_sun / (rho_sun * rho_sun);
416 *derivative = [
417 state[3],
418 state[4],
419 state[5],
420 2.0 * state[4]
421 + state[0]
422 + a1[0]
423 + a2[0]
424 + sun_acceleration[0]
425 + indirect_scale * sun_direction[0],
426 -2.0 * state[3]
427 + state[1]
428 + a1[1]
429 + a2[1]
430 + sun_acceleration[1]
431 + indirect_scale * sun_direction[1],
432 a1[2] + a2[2] + sun_acceleration[2],
433 ];
434 validate_output(Self::NAME, derivative)
435 }
436}
437
438impl DifferentiableDynamicsModel<6, 4> for BcpDynamics {
439 fn jacobians(
440 &self,
441 time: f64,
442 state: &CartesianState,
443 parameters: &[f64; 4],
444 state_jacobian: &mut Matrix6,
445 parameter_jacobian: &mut [[f64; 4]; 6],
446 ) -> Result<()> {
447 self.validate(time, state, parameters)?;
448 let [mu, mu_sun, rho_sun, omega_sun] = *parameters;
449 let (d1, d2) = primary_displacements(state, mu);
450 let (_, gradient1) =
451 point_mass_acceleration_and_gradient(d1, 1.0 - mu, "BCP primary-one distance")?;
452 let (_, gradient2) =
453 point_mass_acceleration_and_gradient(d2, mu, "BCP primary-two distance")?;
454 let (sun_direction, sun_direction_rate, sun_displacement) =
455 sun_geometry(time, state, rho_sun, omega_sun);
456 let (_, sun_gradient) =
457 point_mass_acceleration_and_gradient(sun_displacement, mu_sun, "BCP Sun distance")?;
458 *state_jacobian = rotating_kinematic_jacobian();
459 add_spatial_gradient(state_jacobian, &gradient1);
460 add_spatial_gradient(state_jacobian, &gradient2);
461 add_spatial_gradient(state_jacobian, &sun_gradient);
462
463 *parameter_jacobian = [[0.0; 4]; 6];
464 let mu_derivative1 =
465 moving_mass_parameter_derivative(d1, 1.0 - mu, -1.0, "BCP primary-one distance")?;
466 let mu_derivative2 =
467 moving_mass_parameter_derivative(d2, mu, 1.0, "BCP primary-two distance")?;
468 let sun_radius = radius_squared(sun_displacement, "BCP Sun distance")?.sqrt();
469 let inverse_sun_radius_cubed = 1.0 / sun_radius.powi(3);
470 let inverse_rho_squared = 1.0 / (rho_sun * rho_sun);
471 let inverse_rho_cubed = inverse_rho_squared / rho_sun;
472 for row in 0..POSITION_DIMENSION {
473 parameter_jacobian[row + 3][0] = mu_derivative1[row] + mu_derivative2[row];
474 parameter_jacobian[row + 3][1] = -sun_displacement[row] * inverse_sun_radius_cubed
475 - sun_direction[row] * inverse_rho_squared;
476 let gradient_times_direction = dot_row(&sun_gradient, row, &sun_direction);
477 parameter_jacobian[row + 3][2] =
478 -gradient_times_direction + 2.0 * mu_sun * sun_direction[row] * inverse_rho_cubed;
479 let gradient_times_rate = dot_row(&sun_gradient, row, &sun_direction_rate);
480 parameter_jacobian[row + 3][3] = -rho_sun * time * gradient_times_rate
481 - mu_sun * time * sun_direction_rate[row] * inverse_rho_squared;
482 }
483 validate_jacobians(Self::NAME, state_jacobian, parameter_jacobian)
484 }
485}
486
487fn propagate_stm<M, const P: usize>(
488 model: &M,
489 initial_time: f64,
490 initial_state: CartesianState,
491 final_time: f64,
492 parameters: [f64; P],
493 options: IntegratorOptions,
494) -> Result<SensitivityPropagation<6, 6>>
495where
496 M: DifferentiableDynamicsModel<6, P>,
497{
498 Dop853.propagate_with_sensitivities(
499 model,
500 SensitivityProblem {
501 nominal: InitialValueProblem::new(initial_time, initial_state, final_time, parameters),
502 initial_sensitivities: identity6(),
503 parameter_seeds: [[0.0; 6]; P],
504 },
505 options,
506 )
507}
508
509const fn identity6() -> Matrix6 {
510 let mut matrix = [[0.0; 6]; 6];
511 let mut index = 0;
512 while index < 6 {
513 matrix[index][index] = 1.0;
514 index += 1;
515 }
516 matrix
517}
518
519fn validate_time_state(time: f64, state: &CartesianState) -> Result<()> {
520 ensure_finite("time", time)?;
521 for &value in state {
522 ensure_finite("state", value)?;
523 }
524 Ok(())
525}
526
527fn validate_positive(parameter: &'static str, value: f64) -> Result<()> {
528 ensure_finite(parameter, value)?;
529 if value > 0.0 {
530 Ok(())
531 } else {
532 Err(PykepError::InvalidInput {
533 parameter,
534 reason: "must be greater than zero".into(),
535 })
536 }
537}
538
539fn validate_non_negative(parameter: &'static str, value: f64) -> Result<()> {
540 ensure_finite(parameter, value)?;
541 if value >= 0.0 {
542 Ok(())
543 } else {
544 Err(PykepError::InvalidInput {
545 parameter,
546 reason: "must be non-negative".into(),
547 })
548 }
549}
550
551fn validate_mass_fraction(mu: f64) -> Result<()> {
552 ensure_finite("mu", mu)?;
553 if (0.0..=1.0).contains(&mu) {
554 Ok(())
555 } else {
556 Err(PykepError::InvalidInput {
557 parameter: "mu",
558 reason: "must lie in the closed interval [0, 1]".into(),
559 })
560 }
561}
562
563fn primary_displacements(state: &CartesianState, mu: f64) -> ([f64; 3], [f64; 3]) {
564 (
565 [state[0] + mu, state[1], state[2]],
566 [state[0] + mu - 1.0, state[1], state[2]],
567 )
568}
569
570fn sun_geometry(
571 time: f64,
572 state: &CartesianState,
573 rho_sun: f64,
574 omega_sun: f64,
575) -> ([f64; 3], [f64; 3], [f64; 3]) {
576 let angle = omega_sun * time;
577 let direction = [angle.cos(), angle.sin(), 0.0];
578 let direction_rate = [-angle.sin(), angle.cos(), 0.0];
579 let displacement = [
580 state[0] - rho_sun * direction[0],
581 state[1] - rho_sun * direction[1],
582 state[2],
583 ];
584 (direction, direction_rate, displacement)
585}
586
587fn radius_squared(vector: [f64; 3], operation: &'static str) -> Result<f64> {
588 let squared = vector[0] * vector[0] + vector[1] * vector[1] + vector[2] * vector[2];
589 if squared == 0.0 {
590 return Err(PykepError::SingularGeometry { operation });
591 }
592 if squared.is_finite() {
593 Ok(squared)
594 } else {
595 Err(PykepError::NumericalOverflow { operation })
596 }
597}
598
599fn point_mass_acceleration_and_gradient(
600 displacement: [f64; 3],
601 mass: f64,
602 operation: &'static str,
603) -> Result<([f64; 3], [[f64; 3]; 3])> {
604 let radius_squared = radius_squared(displacement, operation)?;
605 let radius = radius_squared.sqrt();
606 let inverse_radius_cubed = 1.0 / (radius_squared * radius);
607 let inverse_radius_fifth = inverse_radius_cubed / radius_squared;
608 let mut acceleration = [0.0; 3];
609 let mut gradient = [[0.0; 3]; 3];
610 for row in 0..POSITION_DIMENSION {
611 acceleration[row] = -mass * displacement[row] * inverse_radius_cubed;
612 for column in 0..POSITION_DIMENSION {
613 let identity = f64::from(row == column);
614 gradient[row][column] = mass
615 * (3.0 * displacement[row] * displacement[column] * inverse_radius_fifth
616 - identity * inverse_radius_cubed);
617 }
618 }
619 if acceleration
620 .iter()
621 .chain(gradient.iter().flatten())
622 .all(|value| value.is_finite())
623 {
624 Ok((acceleration, gradient))
625 } else {
626 Err(PykepError::NumericalOverflow { operation })
627 }
628}
629
630fn moving_mass_parameter_derivative(
631 displacement: [f64; 3],
632 mass: f64,
633 mass_derivative: f64,
634 operation: &'static str,
635) -> Result<[f64; 3]> {
636 let radius_squared = radius_squared(displacement, operation)?;
637 let radius = radius_squared.sqrt();
638 let inverse_radius_cubed = 1.0 / (radius_squared * radius);
639 let inverse_radius_fifth = inverse_radius_cubed / radius_squared;
640 let mut derivative = [0.0; 3];
641 for row in 0..POSITION_DIMENSION {
642 let displacement_derivative = f64::from(row == 0);
643 derivative[row] = -mass_derivative * displacement[row] * inverse_radius_cubed
644 - mass
645 * (displacement_derivative * inverse_radius_cubed
646 - 3.0 * displacement[row] * displacement[0] * inverse_radius_fifth);
647 }
648 if derivative.iter().all(|value| value.is_finite()) {
649 Ok(derivative)
650 } else {
651 Err(PykepError::NumericalOverflow { operation })
652 }
653}
654
655const fn kinematic_jacobian() -> Matrix6 {
656 let mut matrix = [[0.0; 6]; 6];
657 matrix[0][3] = 1.0;
658 matrix[1][4] = 1.0;
659 matrix[2][5] = 1.0;
660 matrix
661}
662
663const fn rotating_kinematic_jacobian() -> Matrix6 {
664 let mut matrix = kinematic_jacobian();
665 matrix[3][0] = 1.0;
666 matrix[4][1] = 1.0;
667 matrix[3][4] = 2.0;
668 matrix[4][3] = -2.0;
669 matrix
670}
671
672fn set_spatial_gradient(jacobian: &mut Matrix6, gradient: &[[f64; 3]; 3]) {
673 for row in 0..POSITION_DIMENSION {
674 for column in 0..POSITION_DIMENSION {
675 jacobian[row + 3][column] = gradient[row][column];
676 }
677 }
678}
679
680fn add_spatial_gradient(jacobian: &mut Matrix6, gradient: &[[f64; 3]; 3]) {
681 for row in 0..POSITION_DIMENSION {
682 for column in 0..POSITION_DIMENSION {
683 jacobian[row + 3][column] += gradient[row][column];
684 }
685 }
686}
687
688fn dot_row(matrix: &[[f64; 3]; 3], row: usize, vector: &[f64; 3]) -> f64 {
689 matrix[row][0] * vector[0] + matrix[row][1] * vector[1] + matrix[row][2] * vector[2]
690}
691
692fn validate_output(model: &'static str, values: &CartesianState) -> Result<()> {
693 if values.iter().all(|value| value.is_finite()) {
694 Ok(())
695 } else {
696 Err(PykepError::NumericalOverflow { operation: model })
697 }
698}
699
700fn finite_output(model: &'static str, value: f64) -> Result<f64> {
701 if value.is_finite() {
702 Ok(value)
703 } else {
704 Err(PykepError::NumericalOverflow { operation: model })
705 }
706}
707
708fn validate_jacobians<const P: usize>(
709 model: &'static str,
710 state_jacobian: &Matrix6,
711 parameter_jacobian: &[[f64; P]; 6],
712) -> Result<()> {
713 if state_jacobian
714 .iter()
715 .flatten()
716 .chain(parameter_jacobian.iter().flatten())
717 .all(|value| value.is_finite())
718 {
719 Ok(())
720 } else {
721 Err(PykepError::NumericalOverflow { operation: model })
722 }
723}