pykep_core/dynamics/
pontryagin.rs

1// Copyright (c) 2023-2026 Dario Izzo (dario.izzo@gmail.com)
2//                          Advanced Concepts Team, European Space Agency (ESA)
3// Copyright (c) 2026 pykep-rust contributors
4// SPDX-License-Identifier: MPL-2.0
5
6//! Cartesian and modified-equinoctial Pontryagin dynamics.
7//!
8//! This evaluated implementation adapts
9//! `src/ta/pontryagin_cartesian.cpp` and
10//! `src/ta/pontryagin_equinoctial.cpp` from the pinned pykep/kep3 source.
11//! Canonical costate rates use forward-mode differentiation. Full model
12//! Jacobians for propagated sensitivities use fixed-size centered differences
13//! with a relative step of `3e-6`; integrator tolerances do not imply the same
14//! accuracy for those sensitivities.
15
16use core::ops::{Add, Div, Mul, Neg, Sub};
17
18use crate::error::ensure_finite;
19use crate::integration::{DifferentiableDynamicsModel, DynamicsModel};
20use crate::{PykepError, Result};
21
22const PHYSICAL_DIMENSION: usize = 7;
23const AUGMENTED_DIMENSION: usize = 14;
24
25/// Supported indirect optimal-control objectives.
26#[derive(Clone, Copy, Debug, PartialEq, Eq)]
27pub enum Optimality {
28    /// Maximize final mass using the upstream logarithmic throttle barrier.
29    Mass,
30    /// Minimize time with full throttle.
31    Time,
32}
33
34/// Evaluated minimizing control and switching function.
35#[derive(Clone, Copy, Debug, PartialEq)]
36pub struct OptimalControl {
37    /// Optimal throttle in `(0, 1)` for mass optimality and exactly `1` for
38    /// time optimality.
39    pub throttle: f64,
40    /// Optimal Cartesian or RTN thrust direction.
41    pub direction: [f64; 3],
42    /// Switching function `rho`.
43    pub switching_function: f64,
44}
45
46#[derive(Clone, Copy)]
47struct Parameters {
48    mu: f64,
49    maximum_thrust: f64,
50    exhaust_velocity: f64,
51    barrier: Option<f64>,
52    lambda0: f64,
53    optimality: Optimality,
54}
55
56impl Parameters {
57    fn mass(values: &[f64; 5]) -> Result<Self> {
58        let parameters = Self {
59            mu: values[0],
60            maximum_thrust: values[1],
61            exhaust_velocity: values[2],
62            barrier: Some(values[3]),
63            lambda0: values[4],
64            optimality: Optimality::Mass,
65        };
66        parameters.validate()?;
67        Ok(parameters)
68    }
69
70    fn time(values: &[f64; 3]) -> Result<Self> {
71        let parameters = Self {
72            mu: values[0],
73            maximum_thrust: values[1],
74            exhaust_velocity: values[2],
75            barrier: None,
76            lambda0: 1.0,
77            optimality: Optimality::Time,
78        };
79        parameters.validate()?;
80        Ok(parameters)
81    }
82
83    fn validate(self) -> Result<()> {
84        validate_positive("mu", self.mu)?;
85        validate_non_negative("maximum_thrust", self.maximum_thrust)?;
86        validate_positive("exhaust_velocity", self.exhaust_velocity)?;
87        validate_positive("lambda0", self.lambda0)?;
88        if let Some(barrier) = self.barrier {
89            validate_positive("barrier", barrier)?;
90        }
91        Ok(())
92    }
93
94    fn throttle(self, switching_function: f64) -> f64 {
95        match self.optimality {
96            Optimality::Mass => {
97                let barrier = self.barrier.expect("mass mode has a barrier");
98                let root = switching_function.hypot(2.0 * barrier);
99                if switching_function >= 0.0 {
100                    2.0 * barrier / (switching_function + 2.0 * barrier + root)
101                } else {
102                    (root - switching_function) / (root - switching_function + 2.0 * barrier)
103                }
104            }
105            Optimality::Time => 1.0,
106        }
107    }
108
109    fn running_cost(self, throttle: f64) -> f64 {
110        let scale = self.lambda0 * self.maximum_thrust / self.exhaust_velocity;
111        match self.optimality {
112            Optimality::Mass => {
113                let barrier = self.barrier.expect("mass mode has a barrier");
114                scale * (throttle - barrier * (throttle * (1.0 - throttle)).ln())
115            }
116            Optimality::Time => scale,
117        }
118    }
119}
120
121/// Cartesian mass-optimal Pontryagin dynamics.
122///
123/// The 14-state order is
124/// `[x,y,z,vx,vy,vz,m,lx,ly,lz,lvx,lvy,lvz,lm]`. Parameter order is
125/// `[mu, maximum_thrust, exhaust_velocity, barrier, lambda0]`.
126#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
127pub struct CartesianMassOptimal;
128
129/// Cartesian time-optimal Pontryagin dynamics.
130///
131/// State order matches [`CartesianMassOptimal`]. Parameter order is
132/// `[mu, maximum_thrust, exhaust_velocity]`; upstream time optimality uses
133/// full throttle and an implicit `lambda0 = 1`.
134#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
135pub struct CartesianTimeOptimal;
136
137/// Modified-equinoctial mass-optimal Pontryagin dynamics.
138///
139/// The 14-state order is
140/// `[p,f,g,h,k,L,m,lp,lf,lg,lh,lk,lL,lm]`. Parameter order is
141/// `[mu, maximum_thrust, exhaust_velocity, barrier, lambda0]`.
142#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
143pub struct EquinoctialMassOptimal;
144
145/// Modified-equinoctial time-optimal Pontryagin dynamics.
146///
147/// State order matches [`EquinoctialMassOptimal`]. Parameter order is
148/// `[mu, maximum_thrust, exhaust_velocity]`; `lambda0 = 1` is implicit.
149#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
150pub struct EquinoctialTimeOptimal;
151
152macro_rules! implement_model {
153    ($model:ty, $parameter_count:literal, $name:literal, $parameters:expr, $rhs:ident, $positive:expr) => {
154        impl DynamicsModel<AUGMENTED_DIMENSION, $parameter_count> for $model {
155            const NAME: &'static str = $name;
156
157            fn validate(
158                &self,
159                time: f64,
160                state: &[f64; AUGMENTED_DIMENSION],
161                parameters: &[f64; $parameter_count],
162            ) -> Result<()> {
163                validate_augmented(time, state)?;
164                let parameters = $parameters(parameters)?;
165                validate_model_state(state, parameters, stringify!($rhs))
166            }
167
168            fn rhs(
169                &self,
170                time: f64,
171                state: &[f64; AUGMENTED_DIMENSION],
172                parameters: &[f64; $parameter_count],
173                derivative: &mut [f64; AUGMENTED_DIMENSION],
174            ) -> Result<()> {
175                self.validate(time, state, parameters)?;
176                $rhs(state, $parameters(parameters)?, derivative)
177            }
178        }
179
180        impl DifferentiableDynamicsModel<AUGMENTED_DIMENSION, $parameter_count> for $model {
181            fn jacobians(
182                &self,
183                time: f64,
184                state: &[f64; AUGMENTED_DIMENSION],
185                parameters: &[f64; $parameter_count],
186                state_jacobian: &mut [[f64; AUGMENTED_DIMENSION]; AUGMENTED_DIMENSION],
187                parameter_jacobian: &mut [[f64; $parameter_count]; AUGMENTED_DIMENSION],
188            ) -> Result<()> {
189                numerical_jacobians(
190                    self,
191                    time,
192                    state,
193                    parameters,
194                    state_jacobian,
195                    parameter_jacobian,
196                    &[6],
197                    $positive,
198                )
199            }
200        }
201    };
202}
203
204implement_model!(
205    CartesianMassOptimal,
206    5,
207    "Cartesian mass-optimal Pontryagin dynamics",
208    Parameters::mass,
209    cartesian_rhs,
210    &[0, 1, 2, 3, 4]
211);
212implement_model!(
213    CartesianTimeOptimal,
214    3,
215    "Cartesian time-optimal Pontryagin dynamics",
216    Parameters::time,
217    cartesian_rhs,
218    &[0, 1, 2]
219);
220implement_model!(
221    EquinoctialMassOptimal,
222    5,
223    "equinoctial mass-optimal Pontryagin dynamics",
224    Parameters::mass,
225    equinoctial_rhs,
226    &[0, 1, 2, 3, 4]
227);
228implement_model!(
229    EquinoctialTimeOptimal,
230    3,
231    "equinoctial time-optimal Pontryagin dynamics",
232    Parameters::time,
233    equinoctial_rhs,
234    &[0, 1, 2]
235);
236
237/// Evaluates the Cartesian minimizing control in mass-optimal mode.
238///
239/// # Errors
240///
241/// Returns an error for invalid inputs or a zero velocity-primer norm.
242pub fn cartesian_control_mass(state: &[f64; 14], parameters: &[f64; 5]) -> Result<OptimalControl> {
243    let parameters = Parameters::mass(parameters)?;
244    validate_augmented(0.0, state)?;
245    validate_model_state(state, parameters, "cartesian")?;
246    cartesian_control(state, parameters)
247}
248
249/// Evaluates the Cartesian minimizing control in time-optimal mode.
250///
251/// # Errors
252///
253/// Returns an error for invalid inputs or a zero velocity-primer norm.
254pub fn cartesian_control_time(state: &[f64; 14], parameters: &[f64; 3]) -> Result<OptimalControl> {
255    let parameters = Parameters::time(parameters)?;
256    validate_augmented(0.0, state)?;
257    validate_model_state(state, parameters, "cartesian")?;
258    cartesian_control(state, parameters)
259}
260
261/// Evaluates the equinoctial minimizing control in mass-optimal mode.
262///
263/// # Errors
264///
265/// Returns an error for invalid inputs or a zero transformed-primer norm.
266pub fn equinoctial_control_mass(
267    state: &[f64; 14],
268    parameters: &[f64; 5],
269) -> Result<OptimalControl> {
270    let parameters = Parameters::mass(parameters)?;
271    validate_augmented(0.0, state)?;
272    validate_model_state(state, parameters, "equinoctial")?;
273    equinoctial_control(state, parameters)
274}
275
276/// Evaluates the equinoctial minimizing control in time-optimal mode.
277///
278/// # Errors
279///
280/// Returns an error for invalid inputs or a zero transformed-primer norm.
281pub fn equinoctial_control_time(
282    state: &[f64; 14],
283    parameters: &[f64; 3],
284) -> Result<OptimalControl> {
285    let parameters = Parameters::time(parameters)?;
286    validate_augmented(0.0, state)?;
287    validate_model_state(state, parameters, "equinoctial")?;
288    equinoctial_control(state, parameters)
289}
290
291/// Evaluates the minimized Cartesian Hamiltonian in mass-optimal mode.
292///
293/// # Errors
294///
295/// Returns an error under the same conditions as
296/// [`cartesian_control_mass`].
297pub fn cartesian_hamiltonian_mass(state: &[f64; 14], parameters: &[f64; 5]) -> Result<f64> {
298    hamiltonian_cartesian_checked(state, Parameters::mass(parameters)?)
299}
300
301/// Evaluates the minimized Cartesian Hamiltonian in time-optimal mode.
302///
303/// # Errors
304///
305/// Returns an error under the same conditions as
306/// [`cartesian_control_time`].
307pub fn cartesian_hamiltonian_time(state: &[f64; 14], parameters: &[f64; 3]) -> Result<f64> {
308    hamiltonian_cartesian_checked(state, Parameters::time(parameters)?)
309}
310
311/// Evaluates the minimized equinoctial Hamiltonian in mass-optimal mode.
312///
313/// # Errors
314///
315/// Returns an error under the same conditions as
316/// [`equinoctial_control_mass`].
317pub fn equinoctial_hamiltonian_mass(state: &[f64; 14], parameters: &[f64; 5]) -> Result<f64> {
318    hamiltonian_equinoctial_checked(state, Parameters::mass(parameters)?)
319}
320
321/// Evaluates the minimized equinoctial Hamiltonian in time-optimal mode.
322///
323/// # Errors
324///
325/// Returns an error under the same conditions as
326/// [`equinoctial_control_time`].
327pub fn equinoctial_hamiltonian_time(state: &[f64; 14], parameters: &[f64; 3]) -> Result<f64> {
328    hamiltonian_equinoctial_checked(state, Parameters::time(parameters)?)
329}
330
331fn validate_augmented(time: f64, state: &[f64; 14]) -> Result<()> {
332    ensure_finite("time", time)?;
333    for &value in state {
334        ensure_finite("state", value)?;
335    }
336    Ok(())
337}
338
339fn validate_model_state(state: &[f64; 14], parameters: Parameters, model: &str) -> Result<()> {
340    validate_positive("mass", state[6])?;
341    if model == "cartesian" || model == "cartesian_rhs" {
342        let radius_squared = state[0] * state[0] + state[1] * state[1] + state[2] * state[2];
343        if radius_squared == 0.0 {
344            return Err(PykepError::SingularGeometry {
345                operation: "Cartesian Pontryagin radius",
346            });
347        }
348        if !radius_squared.is_finite() {
349            return Err(PykepError::NumericalOverflow {
350                operation: "Cartesian Pontryagin radius",
351            });
352        }
353    } else {
354        validate_positive("semilatus_rectum", state[0])?;
355        let w = 1.0 + state[1] * state[5].cos() + state[2] * state[5].sin();
356        if w == 0.0 {
357            return Err(PykepError::SingularGeometry {
358                operation: "equinoctial Pontryagin radial denominator",
359            });
360        }
361    }
362    let control = if model == "cartesian" || model == "cartesian_rhs" {
363        cartesian_control(state, parameters)
364    } else {
365        equinoctial_control(state, parameters)
366    }?;
367    if control.throttle.is_finite()
368        && control.switching_function.is_finite()
369        && control.direction.iter().all(|value| value.is_finite())
370    {
371        Ok(())
372    } else {
373        Err(PykepError::NumericalOverflow {
374            operation: "Pontryagin minimizing control",
375        })
376    }
377}
378
379fn validate_positive(parameter: &'static str, value: f64) -> Result<()> {
380    ensure_finite(parameter, value)?;
381    if value > 0.0 {
382        Ok(())
383    } else {
384        Err(PykepError::InvalidInput {
385            parameter,
386            reason: "must be greater than zero".into(),
387        })
388    }
389}
390
391fn validate_non_negative(parameter: &'static str, value: f64) -> Result<()> {
392    ensure_finite(parameter, value)?;
393    if value >= 0.0 {
394        Ok(())
395    } else {
396        Err(PykepError::InvalidInput {
397            parameter,
398            reason: "must be greater than or equal to zero".into(),
399        })
400    }
401}
402
403fn cartesian_control(state: &[f64; 14], parameters: Parameters) -> Result<OptimalControl> {
404    let primer = [state[10], state[11], state[12]];
405    let norm = norm(primer);
406    if norm == 0.0 {
407        return Err(PykepError::SingularGeometry {
408            operation: "Cartesian Pontryagin primer norm",
409        });
410    }
411    let direction = [-primer[0] / norm, -primer[1] / norm, -primer[2] / norm];
412    let switching_function = match parameters.optimality {
413        Optimality::Mass => {
414            1.0 - parameters.exhaust_velocity * norm / state[6] / parameters.lambda0
415                - state[13] / parameters.lambda0
416        }
417        Optimality::Time => {
418            -parameters.exhaust_velocity * norm / state[6] / parameters.lambda0
419                - state[13] / parameters.lambda0
420        }
421    };
422    Ok(OptimalControl {
423        throttle: parameters.throttle(switching_function),
424        direction,
425        switching_function,
426    })
427}
428
429fn cartesian_rhs(
430    state: &[f64; 14],
431    parameters: Parameters,
432    derivative: &mut [f64; 14],
433) -> Result<()> {
434    let control = cartesian_control(state, parameters)?;
435    let variables = core::array::from_fn(|index| Dual::variable(state[index], index));
436    let physical = cartesian_physical_dual(variables, parameters, control);
437    let mut hamiltonian = Dual::constant(parameters.running_cost(control.throttle));
438    for index in 0..PHYSICAL_DIMENSION {
439        derivative[index] = physical[index].value;
440        hamiltonian = hamiltonian + Dual::constant(state[index + 7]) * physical[index];
441    }
442    for index in 0..PHYSICAL_DIMENSION {
443        derivative[index + 7] = -hamiltonian.derivative[index];
444    }
445    validate_derivative("Cartesian Pontryagin dynamics", derivative)
446}
447
448fn hamiltonian_cartesian_checked(state: &[f64; 14], parameters: Parameters) -> Result<f64> {
449    validate_augmented(0.0, state)?;
450    validate_model_state(state, parameters, "cartesian")?;
451    let control = cartesian_control(state, parameters)?;
452    let variables = core::array::from_fn(|index| Dual::constant(state[index]));
453    let physical = cartesian_physical_dual(variables, parameters, control);
454    let mut hamiltonian = parameters.running_cost(control.throttle);
455    for index in 0..PHYSICAL_DIMENSION {
456        hamiltonian += state[index + 7] * physical[index].value;
457    }
458    finite("Cartesian Pontryagin Hamiltonian", hamiltonian)
459}
460
461fn cartesian_physical_dual(
462    state: [Dual; 7],
463    parameters: Parameters,
464    control: OptimalControl,
465) -> [Dual; 7] {
466    let radius_squared = state[0] * state[0] + state[1] * state[1] + state[2] * state[2];
467    let gravity_scale = Dual::constant(-parameters.mu) / radius_squared.powf(1.5);
468    let thrust_scale = Dual::constant(parameters.maximum_thrust * control.throttle) / state[6];
469    [
470        state[3],
471        state[4],
472        state[5],
473        gravity_scale * state[0] + thrust_scale * Dual::constant(control.direction[0]),
474        gravity_scale * state[1] + thrust_scale * Dual::constant(control.direction[1]),
475        gravity_scale * state[2] + thrust_scale * Dual::constant(control.direction[2]),
476        Dual::constant(-parameters.maximum_thrust / parameters.exhaust_velocity * control.throttle),
477    ]
478}
479
480fn equinoctial_control(state: &[f64; 14], parameters: Parameters) -> Result<OptimalControl> {
481    let matrix = equinoctial_b_values(state, parameters.mu);
482    let costate = &state[7..13];
483    let primer = core::array::from_fn(|column| {
484        (0..6)
485            .map(|row| matrix[row][column] * costate[row])
486            .sum::<f64>()
487    });
488    let norm = norm(primer);
489    if norm == 0.0 {
490        return Err(PykepError::SingularGeometry {
491            operation: "equinoctial Pontryagin primer norm",
492        });
493    }
494    let direction = [-primer[0] / norm, -primer[1] / norm, -primer[2] / norm];
495    let switching_function = match parameters.optimality {
496        Optimality::Mass => {
497            1.0 - parameters.exhaust_velocity * norm / state[6] / parameters.lambda0
498                - state[13] / parameters.lambda0
499        }
500        Optimality::Time => {
501            -parameters.exhaust_velocity * norm / state[6] / parameters.lambda0
502                - state[13] / parameters.lambda0
503        }
504    };
505    Ok(OptimalControl {
506        throttle: parameters.throttle(switching_function),
507        direction,
508        switching_function,
509    })
510}
511
512fn equinoctial_rhs(
513    state: &[f64; 14],
514    parameters: Parameters,
515    derivative: &mut [f64; 14],
516) -> Result<()> {
517    let control = equinoctial_control(state, parameters)?;
518    let variables = core::array::from_fn(|index| Dual::variable(state[index], index));
519    let physical = equinoctial_physical_dual(variables, parameters, control);
520    let mut hamiltonian = Dual::constant(parameters.running_cost(control.throttle));
521    for index in 0..PHYSICAL_DIMENSION {
522        derivative[index] = physical[index].value;
523        hamiltonian = hamiltonian + Dual::constant(state[index + 7]) * physical[index];
524    }
525    for index in 0..PHYSICAL_DIMENSION {
526        derivative[index + 7] = -hamiltonian.derivative[index];
527    }
528    validate_derivative("equinoctial Pontryagin dynamics", derivative)
529}
530
531fn hamiltonian_equinoctial_checked(state: &[f64; 14], parameters: Parameters) -> Result<f64> {
532    validate_augmented(0.0, state)?;
533    validate_model_state(state, parameters, "equinoctial")?;
534    let control = equinoctial_control(state, parameters)?;
535    let variables = core::array::from_fn(|index| Dual::constant(state[index]));
536    let physical = equinoctial_physical_dual(variables, parameters, control);
537    let mut hamiltonian = parameters.running_cost(control.throttle);
538    for index in 0..PHYSICAL_DIMENSION {
539        hamiltonian += state[index + 7] * physical[index].value;
540    }
541    finite("equinoctial Pontryagin Hamiltonian", hamiltonian)
542}
543
544fn equinoctial_b_values(state: &[f64; 14], mu: f64) -> [[f64; 3]; 6] {
545    let [p, f, g, h, k, longitude, ..] = *state;
546    let sine = longitude.sin();
547    let cosine = longitude.cos();
548    let w = 1.0 + f * cosine + g * sine;
549    let s2 = 1.0 + h * h + k * k;
550    let hsk = h * sine - k * cosine;
551    let scale = (p / mu).sqrt();
552    [
553        [0.0, scale * 2.0 * p / w, 0.0],
554        [
555            scale * sine,
556            scale * ((1.0 + w) * cosine + f) / w,
557            -scale * g * hsk / w,
558        ],
559        [
560            -scale * cosine,
561            scale * ((1.0 + w) * sine + g) / w,
562            scale * f * hsk / w,
563        ],
564        [0.0, 0.0, scale * s2 * cosine / (2.0 * w)],
565        [0.0, 0.0, scale * s2 * sine / (2.0 * w)],
566        [0.0, 0.0, scale * hsk / w],
567    ]
568}
569
570fn equinoctial_physical_dual(
571    state: [Dual; 7],
572    parameters: Parameters,
573    control: OptimalControl,
574) -> [Dual; 7] {
575    let one = Dual::constant(1.0);
576    let sine = state[5].sin();
577    let cosine = state[5].cos();
578    let w = one + state[1] * cosine + state[2] * sine;
579    let s2 = one + state[3] * state[3] + state[4] * state[4];
580    let hsk = state[3] * sine - state[4] * cosine;
581    let scale = (state[0] / Dual::constant(parameters.mu)).sqrt();
582    let matrix = [
583        [
584            Dual::constant(0.0),
585            scale * Dual::constant(2.0) * state[0] / w,
586            Dual::constant(0.0),
587        ],
588        [
589            scale * sine,
590            scale * ((one + w) * cosine + state[1]) / w,
591            -scale * state[2] * hsk / w,
592        ],
593        [
594            -scale * cosine,
595            scale * ((one + w) * sine + state[2]) / w,
596            scale * state[1] * hsk / w,
597        ],
598        [
599            Dual::constant(0.0),
600            Dual::constant(0.0),
601            scale * s2 * cosine / (Dual::constant(2.0) * w),
602        ],
603        [
604            Dual::constant(0.0),
605            Dual::constant(0.0),
606            scale * s2 * sine / (Dual::constant(2.0) * w),
607        ],
608        [Dual::constant(0.0), Dual::constant(0.0), scale * hsk / w],
609    ];
610    let thrust_scale = Dual::constant(parameters.maximum_thrust * control.throttle) / state[6];
611    let mut derivative = [Dual::constant(0.0); 7];
612    for row in 0..6 {
613        derivative[row] = (matrix[row][0] * Dual::constant(control.direction[0])
614            + matrix[row][1] * Dual::constant(control.direction[1])
615            + matrix[row][2] * Dual::constant(control.direction[2]))
616            * thrust_scale;
617    }
618    derivative[5] =
619        derivative[5] + (Dual::constant(parameters.mu) / state[0].powf(3.0)).sqrt() * w * w;
620    derivative[6] = Dual::constant(-parameters.maximum_thrust / parameters.exhaust_velocity)
621        * Dual::constant(control.throttle)
622        * (Dual::constant(-1.0) / state[6] / Dual::constant(1e10)).exp();
623    derivative
624}
625
626fn norm(vector: [f64; 3]) -> f64 {
627    vector.iter().map(|value| value * value).sum::<f64>().sqrt()
628}
629
630fn finite(operation: &'static str, value: f64) -> Result<f64> {
631    if value.is_finite() {
632        Ok(value)
633    } else {
634        Err(PykepError::NumericalOverflow { operation })
635    }
636}
637
638fn validate_derivative(operation: &'static str, derivative: &[f64; 14]) -> Result<()> {
639    if derivative.iter().all(|value| value.is_finite()) {
640        Ok(())
641    } else {
642        Err(PykepError::NumericalOverflow { operation })
643    }
644}
645
646#[allow(clippy::too_many_arguments)]
647fn numerical_jacobians<M, const P: usize>(
648    model: &M,
649    time: f64,
650    state: &[f64; 14],
651    parameters: &[f64; P],
652    state_jacobian: &mut [[f64; 14]; 14],
653    parameter_jacobian: &mut [[f64; P]; 14],
654    positive_state_indices: &[usize],
655    positive_parameter_indices: &[usize],
656) -> Result<()>
657where
658    M: DynamicsModel<14, P>,
659{
660    *state_jacobian = [[0.0; 14]; 14];
661    *parameter_jacobian = [[0.0; P]; 14];
662    for column in 0..14 {
663        let mut step = 3e-6 * state[column].abs().max(1.0);
664        if positive_state_indices.contains(&column) {
665            step = step.min(state[column] * 0.25);
666        }
667        let mut plus = *state;
668        let mut minus = *state;
669        plus[column] += step;
670        minus[column] -= step;
671        let mut rhs_plus = [0.0; 14];
672        let mut rhs_minus = [0.0; 14];
673        model.rhs(time, &plus, parameters, &mut rhs_plus)?;
674        model.rhs(time, &minus, parameters, &mut rhs_minus)?;
675        for row in 0..14 {
676            state_jacobian[row][column] = (rhs_plus[row] - rhs_minus[row]) / (2.0 * step);
677        }
678    }
679    for column in 0..P {
680        let mut step = 3e-6 * parameters[column].abs().max(1.0);
681        if positive_parameter_indices.contains(&column) {
682            step = step.min(parameters[column] * 0.25);
683        }
684        let mut plus = *parameters;
685        let mut minus = *parameters;
686        plus[column] += step;
687        minus[column] -= step;
688        let mut rhs_plus = [0.0; 14];
689        let mut rhs_minus = [0.0; 14];
690        model.rhs(time, state, &plus, &mut rhs_plus)?;
691        model.rhs(time, state, &minus, &mut rhs_minus)?;
692        for row in 0..14 {
693            parameter_jacobian[row][column] = (rhs_plus[row] - rhs_minus[row]) / (2.0 * step);
694        }
695    }
696    Ok(())
697}
698
699#[derive(Clone, Copy, Debug)]
700struct Dual {
701    value: f64,
702    derivative: [f64; PHYSICAL_DIMENSION],
703}
704
705impl Dual {
706    const fn constant(value: f64) -> Self {
707        Self {
708            value,
709            derivative: [0.0; PHYSICAL_DIMENSION],
710        }
711    }
712
713    fn variable(value: f64, index: usize) -> Self {
714        let mut derivative = [0.0; PHYSICAL_DIMENSION];
715        derivative[index] = 1.0;
716        Self { value, derivative }
717    }
718
719    fn sqrt(self) -> Self {
720        self.powf(0.5)
721    }
722
723    fn powf(self, exponent: f64) -> Self {
724        let value = self.value.powf(exponent);
725        let scale = exponent * self.value.powf(exponent - 1.0);
726        Self {
727            value,
728            derivative: self.derivative.map(|item| item * scale),
729        }
730    }
731
732    fn exp(self) -> Self {
733        let value = self.value.exp();
734        Self {
735            value,
736            derivative: self.derivative.map(|item| item * value),
737        }
738    }
739
740    fn sin(self) -> Self {
741        let value = self.value.sin();
742        let scale = self.value.cos();
743        Self {
744            value,
745            derivative: self.derivative.map(|item| item * scale),
746        }
747    }
748
749    fn cos(self) -> Self {
750        let value = self.value.cos();
751        let scale = -self.value.sin();
752        Self {
753            value,
754            derivative: self.derivative.map(|item| item * scale),
755        }
756    }
757}
758
759macro_rules! dual_binary {
760    ($trait:ident, $method:ident, $value:expr, $derivative:expr) => {
761        impl $trait for Dual {
762            type Output = Self;
763
764            fn $method(self, right: Self) -> Self {
765                let mut derivative = [0.0; PHYSICAL_DIMENSION];
766                for (index, item) in derivative.iter_mut().enumerate() {
767                    *item = $derivative(self, right, index);
768                }
769                Self {
770                    value: $value(self, right),
771                    derivative,
772                }
773            }
774        }
775    };
776}
777
778dual_binary!(
779    Add,
780    add,
781    |left: Dual, right: Dual| left.value + right.value,
782    |left: Dual, right: Dual, index| left.derivative[index] + right.derivative[index]
783);
784dual_binary!(
785    Sub,
786    sub,
787    |left: Dual, right: Dual| left.value - right.value,
788    |left: Dual, right: Dual, index| left.derivative[index] - right.derivative[index]
789);
790dual_binary!(
791    Mul,
792    mul,
793    |left: Dual, right: Dual| left.value * right.value,
794    |left: Dual, right: Dual, index| left.derivative[index] * right.value
795        + left.value * right.derivative[index]
796);
797dual_binary!(
798    Div,
799    div,
800    |left: Dual, right: Dual| left.value / right.value,
801    |left: Dual, right: Dual, index| {
802        (left.derivative[index] * right.value - left.value * right.derivative[index])
803            / (right.value * right.value)
804    }
805);
806
807impl Neg for Dual {
808    type Output = Self;
809
810    fn neg(self) -> Self {
811        Self {
812            value: -self.value,
813            derivative: self.derivative.map(|item| -item),
814        }
815    }
816}
pykep_core/dynamics/pontryagin.rs

pykep_core/dynamics/
pontryagin.rs
