pykep_core/
integration.rs

1// Copyright (c) 2026 pykep-rust contributors
2// SPDX-License-Identifier: MPL-2.0
3
4//! Adaptive integration for evaluated astrodynamics models.
5//!
6//! This module deliberately exposes pykep-owned model and result types rather
7//! than the implementation types of the selected solver crate. The numerical
8//! backend is DOP853: an explicit Runge-Kutta method of order 8 with embedded
9//! error estimates and seventh-order dense output.
10//!
11//! ```
12//! use pykep_core::integration::{
13//!     Dop853, DynamicsModel, InitialValueProblem, IntegratorOptions,
14//! };
15//!
16//! struct Constant;
17//! impl DynamicsModel<1, 0> for Constant {
18//!     const NAME: &'static str = "constant";
19//!
20//!     fn rhs(
21//!         &self,
22//!         _time: f64,
23//!         _state: &[f64; 1],
24//!         _parameters: &[f64; 0],
25//!         derivative: &mut [f64; 1],
26//!     ) -> pykep_core::Result<()> {
27//!         *derivative = [2.0];
28//!         Ok(())
29//!     }
30//! }
31//!
32//! let result = Dop853.propagate(
33//!     &Constant,
34//!     InitialValueProblem::new(0.0, [1.0], 0.5, []),
35//!     IntegratorOptions::default(),
36//! )?;
37//! assert!((result.state[0] - 2.0).abs() < 1e-12);
38//! # Ok::<(), pykep_core::PykepError>(())
39//! ```
40
41use core::cell::RefCell;
42
43use differential_equations::control::ControlFlag;
44use differential_equations::interpolate::Interpolation;
45use differential_equations::ode::{ODE, solve_ode};
46use differential_equations::prelude::ExplicitRungeKutta;
47use differential_equations::solout::{
48    CrossingDirection, DefaultSolout, Event as BackendEvent, EventConfig, Solout, TEvalSolout,
49};
50use differential_equations::solution::Solution;
51use differential_equations::tolerance::Tolerance as BackendTolerance;
52use differential_equations::traits::State as BackendState;
53
54use crate::error::ensure_finite;
55use crate::{PykepError, Result};
56
57/// An evaluated first-order dynamics model with `N` states and `P` parameters.
58///
59/// Implementations must not allocate in [`Self::rhs`]. The integrator validates
60/// dimensions and finite inputs before calling the model.
61pub trait DynamicsModel<const N: usize, const P: usize> {
62    /// Stable model name used in diagnostics.
63    const NAME: &'static str;
64
65    /// Validates state and parameter domains before a propagation.
66    ///
67    /// The default accepts every finite state and parameter vector.
68    ///
69    /// # Errors
70    ///
71    /// Returns an error when the initial state or parameters are outside the
72    /// model's physical domain.
73    fn validate(&self, _time: f64, _state: &[f64; N], _parameters: &[f64; P]) -> Result<()> {
74        Ok(())
75    }
76
77    /// Evaluates `dstate/dt` into caller-owned storage.
78    ///
79    /// # Errors
80    ///
81    /// Returns an error when the dynamics are singular or cannot be evaluated.
82    fn rhs(
83        &self,
84        time: f64,
85        state: &[f64; N],
86        parameters: &[f64; P],
87        derivative: &mut [f64; N],
88    ) -> Result<()>;
89}
90
91/// A dynamics model that provides the Jacobians required by sensitivities.
92pub trait DifferentiableDynamicsModel<const N: usize, const P: usize>: DynamicsModel<N, P> {
93    /// Evaluates `df/dstate` and `df/dparameters` in row-major order.
94    ///
95    /// # Errors
96    ///
97    /// Returns an error under the same conditions as [`DynamicsModel::rhs`].
98    fn jacobians(
99        &self,
100        time: f64,
101        state: &[f64; N],
102        parameters: &[f64; P],
103        state_jacobian: &mut [[f64; N]; N],
104        parameter_jacobian: &mut [[f64; P]; N],
105    ) -> Result<()>;
106}
107
108/// Direction in which an event function must cross zero.
109#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
110pub enum EventDirection {
111    /// Detect either crossing direction.
112    #[default]
113    Either,
114    /// Detect a transition from negative to non-negative.
115    Increasing,
116    /// Detect a transition from positive to non-positive.
117    Decreasing,
118}
119
120/// A scalar zero-crossing condition for a fixed-size state.
121pub trait Event<const N: usize> {
122    /// Crossing direction to detect.
123    fn direction(&self) -> EventDirection {
124        EventDirection::Either
125    }
126
127    /// Evaluates the scalar event function.
128    fn value(&self, time: f64, state: &[f64; N]) -> f64;
129}
130
131/// Scalar error tolerances and safety limits for adaptive integration.
132#[derive(Clone, Copy, Debug, PartialEq)]
133pub struct IntegratorOptions {
134    /// Relative local-error tolerance.
135    pub relative_tolerance: f64,
136    /// Absolute local-error tolerance.
137    pub absolute_tolerance: f64,
138    /// Optional initial step magnitude. The backend estimates it when absent.
139    pub initial_step: Option<f64>,
140    /// Optional maximum step magnitude.
141    pub maximum_step: Option<f64>,
142    /// Maximum accepted plus rejected steps.
143    pub maximum_steps: usize,
144    /// Maximum consecutive rejected steps.
145    pub maximum_rejections: usize,
146}
147
148impl Default for IntegratorOptions {
149    fn default() -> Self {
150        Self {
151            relative_tolerance: 1e-12,
152            absolute_tolerance: 1e-12,
153            initial_step: None,
154            maximum_step: None,
155            maximum_steps: 100_000,
156            maximum_rejections: 100,
157        }
158    }
159}
160
161impl IntegratorOptions {
162    pub(crate) fn validate(self) -> Result<()> {
163        validate_positive_finite("relative_tolerance", self.relative_tolerance)?;
164        validate_positive_finite("absolute_tolerance", self.absolute_tolerance)?;
165        if let Some(step) = self.initial_step {
166            validate_positive_finite("initial_step", step)?;
167        }
168        if let Some(step) = self.maximum_step {
169            validate_positive_finite("maximum_step", step)?;
170        }
171        if self.maximum_steps == 0 {
172            return Err(PykepError::InvalidInput {
173                parameter: "maximum_steps",
174                reason: "must be greater than zero".into(),
175            });
176        }
177        if self.maximum_rejections == 0 {
178            return Err(PykepError::InvalidInput {
179                parameter: "maximum_rejections",
180                reason: "must be greater than zero".into(),
181            });
182        }
183        Ok(())
184    }
185}
186
187fn validate_positive_finite(parameter: &'static str, value: f64) -> Result<()> {
188    ensure_finite(parameter, value)?;
189    if value > 0.0 {
190        Ok(())
191    } else {
192        Err(PykepError::InvalidInput {
193            parameter,
194            reason: "must be greater than zero".into(),
195        })
196    }
197}
198
199/// Initial-value problem for a fixed-size model.
200#[derive(Clone, Copy, Debug, PartialEq)]
201pub struct InitialValueProblem<const N: usize, const P: usize> {
202    /// Initial integration time.
203    pub initial_time: f64,
204    /// State at [`Self::initial_time`].
205    pub initial_state: [f64; N],
206    /// Requested final integration time.
207    pub final_time: f64,
208    /// Constant model parameters for this propagation.
209    pub parameters: [f64; P],
210}
211
212impl<const N: usize, const P: usize> InitialValueProblem<N, P> {
213    /// Creates an initial-value problem.
214    pub const fn new(
215        initial_time: f64,
216        initial_state: [f64; N],
217        final_time: f64,
218        parameters: [f64; P],
219    ) -> Self {
220        Self {
221            initial_time,
222            initial_state,
223            final_time,
224            parameters,
225        }
226    }
227}
228
229/// Initial-value problem augmented with arbitrary sensitivity seed directions.
230#[derive(Clone, Copy, Debug, PartialEq)]
231pub struct SensitivityProblem<const N: usize, const P: usize, const W: usize> {
232    /// Nominal initial-value problem.
233    pub nominal: InitialValueProblem<N, P>,
234    /// Initial `dstate/dseed` matrix.
235    pub initial_sensitivities: [[f64; W]; N],
236    /// Constant `dparameters/dseed` matrix.
237    pub parameter_seeds: [[f64; W]; P],
238}
239
240/// Work counters reported by a completed propagation.
241#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
242pub struct IntegrationStats {
243    /// Right-hand-side evaluations.
244    pub rhs_evaluations: usize,
245    /// Accepted integration steps.
246    pub accepted_steps: usize,
247    /// Rejected integration steps.
248    pub rejected_steps: usize,
249}
250
251/// Reason that integration stopped.
252#[derive(Clone, Copy, Debug, PartialEq, Eq)]
253pub enum Termination {
254    /// The requested final time was reached.
255    FinalTime,
256    /// A terminal event was located by dense interpolation.
257    Event,
258}
259
260/// Final state and diagnostics from a propagation.
261#[derive(Clone, Debug, PartialEq)]
262pub struct Propagation<const N: usize> {
263    /// Time at which integration stopped.
264    pub time: f64,
265    /// State at [`Self::time`].
266    pub state: [f64; N],
267    /// Integration work counters.
268    pub stats: IntegrationStats,
269    /// Reason that integration stopped.
270    pub termination: Termination,
271}
272
273/// States evaluated at caller-requested times using dense output.
274#[derive(Clone, Debug, PartialEq)]
275pub struct DenseTrajectory<const N: usize> {
276    /// Evaluation times in propagation order.
277    pub times: Vec<f64>,
278    /// States corresponding one-to-one with [`Self::times`].
279    pub states: Vec<[f64; N]>,
280    /// Integration work counters.
281    pub stats: IntegrationStats,
282}
283
284/// Final state and first-order sensitivities.
285///
286/// `sensitivities[i][j]` is the derivative of state component `i` with
287/// respect to seed direction `j`.
288#[derive(Clone, Debug, PartialEq)]
289pub struct SensitivityPropagation<const N: usize, const W: usize> {
290    /// Time at the end of propagation.
291    pub time: f64,
292    /// Nominal propagated state.
293    pub state: [f64; N],
294    /// Row-major state sensitivity matrix.
295    pub sensitivities: [[f64; W]; N],
296    /// Integration work counters.
297    pub stats: IntegrationStats,
298}
299
300/// Pure-Rust adaptive DOP853 integration facade.
301#[derive(Clone, Copy, Debug, Default)]
302pub struct Dop853;
303
304impl Dop853 {
305    /// Propagates a model from `initial_time` to `final_time`.
306    ///
307    /// No heap allocation occurs per accepted step for fixed-size states.
308    ///
309    /// # Errors
310    ///
311    /// Returns an error for invalid options, non-finite inputs, model failures,
312    /// step-size underflow, stiffness detection, or exhausted step limits.
313    pub fn propagate<M, const N: usize, const P: usize>(
314        &self,
315        model: &M,
316        problem: InitialValueProblem<N, P>,
317        options: IntegratorOptions,
318    ) -> Result<Propagation<N>>
319    where
320        M: DynamicsModel<N, P>,
321    {
322        validate_problem(model, &problem, options)?;
323        if problem.initial_time == problem.final_time {
324            return Ok(Propagation {
325                time: problem.initial_time,
326                state: problem.initial_state,
327                stats: IntegrationStats::default(),
328                termination: Termination::FinalTime,
329            });
330        }
331
332        let adapter = ModelAdapter::new(model, &problem.parameters);
333        let mut solver = configured_solver(options);
334        let mut output = FinalState::default();
335        let solution = solve_ode(
336            &mut solver,
337            &adapter,
338            problem.initial_time,
339            problem.final_time,
340            &problem.initial_state,
341            &mut output,
342        );
343        finish_model_error(&adapter)?;
344        let solution = solution.map_err(|error| integration_error(M::NAME, &error))?;
345        let (time, state) = output.value.ok_or_else(|| PykepError::IntegrationFailure {
346            model: M::NAME,
347            reason: "backend returned no final state".into(),
348        })?;
349        ensure_finite_state(M::NAME, &state)?;
350
351        Ok(Propagation {
352            time,
353            state,
354            stats: stats(&solution),
355            termination: Termination::FinalTime,
356        })
357    }
358
359    /// Propagates and evaluates the dense solution at requested times.
360    ///
361    /// `evaluation_times` must be strictly monotone in the propagation
362    /// direction and lie in the closed integration interval.
363    ///
364    /// # Errors
365    ///
366    /// Returns an error for an invalid grid or any error described by
367    /// [`Self::propagate`].
368    pub fn propagate_dense<M, const N: usize, const P: usize>(
369        &self,
370        model: &M,
371        problem: InitialValueProblem<N, P>,
372        evaluation_times: &[f64],
373        options: IntegratorOptions,
374    ) -> Result<DenseTrajectory<N>>
375    where
376        M: DynamicsModel<N, P>,
377    {
378        validate_problem(model, &problem, options)?;
379        validate_evaluation_times(problem.initial_time, problem.final_time, evaluation_times)?;
380        if problem.initial_time == problem.final_time {
381            return Ok(DenseTrajectory {
382                times: vec![problem.initial_time],
383                states: vec![problem.initial_state],
384                stats: IntegrationStats::default(),
385            });
386        }
387
388        let adapter = ModelAdapter::new(model, &problem.parameters);
389        let mut solver = configured_solver(options);
390        let mut output =
391            TEvalSolout::new(evaluation_times, problem.initial_time, problem.final_time);
392        let solution = solve_ode(
393            &mut solver,
394            &adapter,
395            problem.initial_time,
396            problem.final_time,
397            &problem.initial_state,
398            &mut output,
399        );
400        finish_model_error(&adapter)?;
401        let solution = solution.map_err(|error| integration_error(M::NAME, &error))?;
402        for state in &solution.y {
403            ensure_finite_state(M::NAME, state)?;
404        }
405        Ok(DenseTrajectory {
406            times: solution.t.clone(),
407            states: solution.y.clone(),
408            stats: stats(&solution),
409        })
410    }
411
412    /// Propagates until a terminal zero-crossing or the requested final time.
413    ///
414    /// # Errors
415    ///
416    /// Returns an error for a non-finite event value or any error described by
417    /// [`Self::propagate`].
418    pub fn propagate_until_event<M, E, const N: usize, const P: usize>(
419        &self,
420        model: &M,
421        event: &E,
422        problem: InitialValueProblem<N, P>,
423        options: IntegratorOptions,
424    ) -> Result<Propagation<N>>
425    where
426        M: DynamicsModel<N, P>,
427        E: Event<N>,
428    {
429        validate_problem(model, &problem, options)?;
430        if problem.initial_time == problem.final_time {
431            return Ok(Propagation {
432                time: problem.initial_time,
433                state: problem.initial_state,
434                stats: IntegrationStats::default(),
435                termination: Termination::FinalTime,
436            });
437        }
438
439        let adapter = ModelAdapter::new(model, &problem.parameters);
440        let event_adapter = EventAdapter::new(event);
441        let mut solver = configured_solver(options);
442        let mut output = differential_equations::solout::EventWrappedSolout::new(
443            DefaultSolout::new(),
444            &event_adapter,
445            problem.initial_time,
446            problem.final_time,
447        );
448        let solution = solve_ode(
449            &mut solver,
450            &adapter,
451            problem.initial_time,
452            problem.final_time,
453            &problem.initial_state,
454            &mut output,
455        );
456        finish_model_error(&adapter)?;
457        if event_adapter.non_finite_value() {
458            return Err(PykepError::IntegrationFailure {
459                model: M::NAME,
460                reason: "event function returned a non-finite value".into(),
461            });
462        }
463        let solution = solution.map_err(|error| integration_error(M::NAME, &error))?;
464        let (&time, &state) = solution
465            .last()
466            .map_err(|_| PykepError::IntegrationFailure {
467                model: M::NAME,
468                reason: "backend returned no event or final state".into(),
469            })?;
470        ensure_finite_state(M::NAME, &state)?;
471        let termination = if matches!(
472            solution.status,
473            differential_equations::status::Status::Interrupted
474        ) {
475            Termination::Event
476        } else {
477            Termination::FinalTime
478        };
479        Ok(Propagation {
480            time,
481            state,
482            stats: stats(&solution),
483            termination,
484        })
485    }
486
487    /// Propagates a nominal state and arbitrary first-order seed directions.
488    ///
489    /// The initial sensitivity matrix contains `dstate0/dseed`; the parameter
490    /// seed matrix contains `dparameters/dseed`. This supports STMs, parameter
491    /// sensitivities, and the state-plus-control variations required by ZOH
492    /// legs without solver-generated expression graphs.
493    ///
494    /// # Errors
495    ///
496    /// Returns an error for `W == 0`, invalid seeds, Jacobian failures, or any
497    /// error described by [`Self::propagate`].
498    pub fn propagate_with_sensitivities<M, const N: usize, const P: usize, const W: usize>(
499        &self,
500        model: &M,
501        problem: SensitivityProblem<N, P, W>,
502        options: IntegratorOptions,
503    ) -> Result<SensitivityPropagation<N, W>>
504    where
505        M: DifferentiableDynamicsModel<N, P>,
506    {
507        validate_problem(model, &problem.nominal, options)?;
508        if W == 0 {
509            return Err(PykepError::InvalidInput {
510                parameter: "sensitivity_width",
511                reason: "must be greater than zero".into(),
512            });
513        }
514        ensure_finite_matrix("initial_sensitivities", &problem.initial_sensitivities)?;
515        ensure_finite_matrix("parameter_seeds", &problem.parameter_seeds)?;
516        if problem.nominal.initial_time == problem.nominal.final_time {
517            return Ok(SensitivityPropagation {
518                time: problem.nominal.initial_time,
519                state: problem.nominal.initial_state,
520                sensitivities: problem.initial_sensitivities,
521                stats: IntegrationStats::default(),
522            });
523        }
524
525        let augmented = AugmentedModel {
526            model,
527            parameters: &problem.nominal.parameters,
528            parameter_seeds: &problem.parameter_seeds,
529            error: RefCell::new(None),
530        };
531        let initial = AugmentedState {
532            state: problem.nominal.initial_state,
533            sensitivities: problem.initial_sensitivities,
534        };
535        let mut solver = configured_solver(options);
536        let mut output = FinalState::default();
537        let solution = solve_ode(
538            &mut solver,
539            &augmented,
540            problem.nominal.initial_time,
541            problem.nominal.final_time,
542            &initial,
543            &mut output,
544        );
545        if let Some(error) = augmented.error.borrow_mut().take() {
546            return Err(error);
547        }
548        let solution = solution.map_err(|error| integration_error(M::NAME, &error))?;
549        let (time, final_state) = output.value.ok_or_else(|| PykepError::IntegrationFailure {
550            model: M::NAME,
551            reason: "backend returned no sensitivity state".into(),
552        })?;
553        ensure_finite_state(M::NAME, &final_state.state)?;
554        ensure_finite_matrix("sensitivities", &final_state.sensitivities)?;
555        Ok(SensitivityPropagation {
556            time,
557            state: final_state.state,
558            sensitivities: final_state.sensitivities,
559            stats: stats(&solution),
560        })
561    }
562}
563
564fn configured_solver<Y>(
565    options: IntegratorOptions,
566) -> differential_equations::methods::ExplicitRungeKutta<
567    differential_equations::methods::Ordinary,
568    differential_equations::methods::DormandPrince,
569    f64,
570    Y,
571    8,
572    12,
573    16,
574>
575where
576    Y: BackendState<f64>,
577{
578    ExplicitRungeKutta::dop853()
579        .rtol(options.relative_tolerance)
580        .atol(options.absolute_tolerance)
581        .h0(options.initial_step.unwrap_or(0.0))
582        .h_max(options.maximum_step.unwrap_or(f64::INFINITY))
583        .max_steps(options.maximum_steps)
584        .max_rejects(options.maximum_rejections)
585}
586
587fn validate_problem<M, const N: usize, const P: usize>(
588    model: &M,
589    problem: &InitialValueProblem<N, P>,
590    options: IntegratorOptions,
591) -> Result<()>
592where
593    M: DynamicsModel<N, P>,
594{
595    if N == 0 {
596        return Err(PykepError::InvalidInput {
597            parameter: "state_dimension",
598            reason: "must be greater than zero".into(),
599        });
600    }
601    ensure_finite("initial_time", problem.initial_time)?;
602    ensure_finite("final_time", problem.final_time)?;
603    ensure_finite_matrix("initial_state", &problem.initial_state)?;
604    ensure_finite_matrix("parameters", &problem.parameters)?;
605    options.validate()?;
606    model.validate(
607        problem.initial_time,
608        &problem.initial_state,
609        &problem.parameters,
610    )
611}
612
613fn validate_evaluation_times(initial: f64, final_time: f64, times: &[f64]) -> Result<()> {
614    if times.is_empty() {
615        return Err(PykepError::InvalidInput {
616            parameter: "evaluation_times",
617            reason: "must contain at least one time".into(),
618        });
619    }
620    if initial == final_time {
621        if times.len() == 1 && times[0] == initial {
622            return Ok(());
623        }
624        return Err(PykepError::InvalidInput {
625            parameter: "evaluation_times",
626            reason: "a zero-duration propagation accepts only the initial time".into(),
627        });
628    }
629    let direction = (final_time - initial).signum();
630    let lower = initial.min(final_time);
631    let upper = initial.max(final_time);
632    for (index, &time) in times.iter().enumerate() {
633        ensure_finite("evaluation_time", time)?;
634        if !(lower..=upper).contains(&time) {
635            return Err(PykepError::InvalidInput {
636                parameter: "evaluation_times",
637                reason: "all times must lie in the closed propagation interval".into(),
638            });
639        }
640        if index > 0 && (time - times[index - 1]) * direction <= 0.0 {
641            return Err(PykepError::InvalidInput {
642                parameter: "evaluation_times",
643                reason: "times must be strictly monotone in the propagation direction".into(),
644            });
645        }
646    }
647    Ok(())
648}
649
650fn ensure_finite_matrix<const R: usize, T>(parameter: &'static str, values: &[T; R]) -> Result<()>
651where
652    T: FiniteValues,
653{
654    for value in values {
655        if !value.all_finite() {
656            return Err(PykepError::NonFiniteInput { parameter });
657        }
658    }
659    Ok(())
660}
661
662trait FiniteValues {
663    fn all_finite(&self) -> bool;
664}
665
666impl FiniteValues for f64 {
667    fn all_finite(&self) -> bool {
668        self.is_finite()
669    }
670}
671
672impl<const N: usize> FiniteValues for [f64; N] {
673    fn all_finite(&self) -> bool {
674        self.iter().all(|value| value.is_finite())
675    }
676}
677
678fn ensure_finite_state<const N: usize>(model: &'static str, state: &[f64; N]) -> Result<()> {
679    if state.iter().all(|value| value.is_finite()) {
680        Ok(())
681    } else {
682        Err(PykepError::IntegrationFailure {
683            model,
684            reason: "propagation produced a non-finite state".into(),
685        })
686    }
687}
688
689fn integration_error<T: core::fmt::Debug>(model: &'static str, error: &T) -> PykepError {
690    PykepError::IntegrationFailure {
691        model,
692        reason: format!("{error:?}"),
693    }
694}
695
696fn stats<T, Y>(solution: &Solution<T, Y>) -> IntegrationStats
697where
698    T: differential_equations::traits::Real,
699    Y: BackendState<T>,
700{
701    IntegrationStats {
702        rhs_evaluations: solution.evals.function,
703        accepted_steps: solution.steps.accepted,
704        rejected_steps: solution.steps.rejected,
705    }
706}
707
708struct ModelAdapter<'a, M, const N: usize, const P: usize> {
709    model: &'a M,
710    parameters: &'a [f64; P],
711    error: RefCell<Option<PykepError>>,
712}
713
714impl<'a, M, const N: usize, const P: usize> ModelAdapter<'a, M, N, P> {
715    fn new(model: &'a M, parameters: &'a [f64; P]) -> Self {
716        Self {
717            model,
718            parameters,
719            error: RefCell::new(None),
720        }
721    }
722}
723
724impl<M, const N: usize, const P: usize> ODE<f64, [f64; N]> for ModelAdapter<'_, M, N, P>
725where
726    M: DynamicsModel<N, P>,
727{
728    fn diff(&self, time: f64, state: &[f64; N], derivative: &mut [f64; N]) {
729        if self.error.borrow().is_some() {
730            derivative.fill(f64::NAN);
731            return;
732        }
733        if let Err(error) = self.model.rhs(time, state, self.parameters, derivative) {
734            *self.error.borrow_mut() = Some(error);
735            derivative.fill(f64::NAN);
736        } else if derivative.iter().any(|value| !value.is_finite()) {
737            *self.error.borrow_mut() = Some(PykepError::IntegrationFailure {
738                model: M::NAME,
739                reason: "right-hand side returned a non-finite derivative".into(),
740            });
741            derivative.fill(f64::NAN);
742        }
743    }
744}
745
746fn finish_model_error<M, const N: usize, const P: usize>(
747    adapter: &ModelAdapter<'_, M, N, P>,
748) -> Result<()> {
749    match adapter.error.borrow_mut().take() {
750        Some(error) => Err(error),
751        None => Ok(()),
752    }
753}
754
755#[derive(Clone, Debug)]
756struct FinalState<Y> {
757    value: Option<(f64, Y)>,
758}
759
760impl<Y> Default for FinalState<Y> {
761    fn default() -> Self {
762        Self { value: None }
763    }
764}
765
766impl<Y> Solout<f64, Y> for FinalState<Y>
767where
768    Y: BackendState<f64>,
769{
770    fn solout<I>(
771        &mut self,
772        time: f64,
773        _previous_time: f64,
774        state: &Y,
775        _previous_state: &Y,
776        _interpolator: &mut I,
777        _solution: &mut Solution<f64, Y>,
778    ) -> ControlFlag<f64, Y>
779    where
780        I: Interpolation<f64, Y> + ?Sized,
781    {
782        self.value = Some((time, state.clone()));
783        ControlFlag::Continue
784    }
785}
786
787struct EventAdapter<'a, E> {
788    event: &'a E,
789    non_finite: core::cell::Cell<bool>,
790}
791
792impl<'a, E> EventAdapter<'a, E> {
793    fn new(event: &'a E) -> Self {
794        Self {
795            event,
796            non_finite: core::cell::Cell::new(false),
797        }
798    }
799
800    fn non_finite_value(&self) -> bool {
801        self.non_finite.get()
802    }
803}
804
805impl<E, const N: usize> BackendEvent<f64, [f64; N]> for EventAdapter<'_, E>
806where
807    E: Event<N>,
808{
809    fn config(&self) -> EventConfig {
810        let direction = match self.event.direction() {
811            EventDirection::Either => CrossingDirection::Both,
812            EventDirection::Increasing => CrossingDirection::Positive,
813            EventDirection::Decreasing => CrossingDirection::Negative,
814        };
815        EventConfig::new(direction, Some(1))
816    }
817
818    fn event(&self, time: f64, state: &[f64; N]) -> f64 {
819        let value = self.event.value(time, state);
820        if value.is_finite() {
821            value
822        } else {
823            self.non_finite.set(true);
824            f64::NAN
825        }
826    }
827}
828
829#[derive(Clone, Debug, PartialEq)]
830struct AugmentedState<const N: usize, const W: usize> {
831    state: [f64; N],
832    sensitivities: [[f64; W]; N],
833}
834
835impl<const N: usize, const W: usize> AugmentedState<N, W> {
836    fn for_each(&self, mut operation: impl FnMut(usize, f64)) {
837        for (index, &value) in self.state.iter().enumerate() {
838            operation(index, value);
839        }
840        let mut index = N;
841        for row in &self.sensitivities {
842            for &value in row {
843                operation(index, value);
844                index += 1;
845            }
846        }
847    }
848}
849
850impl<const N: usize, const W: usize> BackendState<f64> for AugmentedState<N, W> {
851    fn len(&self) -> usize {
852        N + N * W
853    }
854
855    fn get_component(&self, index: usize) -> f64 {
856        if index < N {
857            self.state[index]
858        } else {
859            let flat = index - N;
860            self.sensitivities[flat / W][flat % W]
861        }
862    }
863
864    fn set_component(&mut self, index: usize, value: f64) {
865        if index < N {
866            self.state[index] = value;
867        } else {
868            let flat = index - N;
869            self.sensitivities[flat / W][flat % W] = value;
870        }
871    }
872
873    fn map_components_mut<F>(&mut self, mut operation: F)
874    where
875        F: FnMut(usize, &mut f64),
876    {
877        for (index, value) in self.state.iter_mut().enumerate() {
878            operation(index, value);
879        }
880        let mut index = N;
881        for row in &mut self.sensitivities {
882            for value in row {
883                operation(index, value);
884                index += 1;
885            }
886        }
887    }
888
889    fn zeros_like(&self) -> Self {
890        Self::zeros()
891    }
892
893    fn zeros() -> Self {
894        Self {
895            state: [0.0; N],
896            sensitivities: [[0.0; W]; N],
897        }
898    }
899
900    fn mul_add_assign(&mut self, alpha: f64, other: &Self) {
901        self.map_components_mut(|index, value| {
902            *value += alpha * other.get_component(index);
903        });
904    }
905
906    fn scale_mut(&mut self, alpha: f64) {
907        self.map_components_mut(|_, value| *value *= alpha);
908    }
909
910    fn norm_squared(&self) -> f64 {
911        let mut value = 0.0;
912        self.for_each(|_, component| value += component * component);
913        value
914    }
915
916    fn diff_norm_squared(&self, other: &Self) -> f64 {
917        let mut value = 0.0;
918        self.for_each(|index, component| {
919            let difference = component - other.get_component(index);
920            value += difference * difference;
921        });
922        value
923    }
924
925    fn error_norm(
926        &self,
927        new_state: &Self,
928        error: &Self,
929        absolute: &BackendTolerance<f64>,
930        relative: &BackendTolerance<f64>,
931    ) -> f64 {
932        let mut value = 0.0;
933        self.for_each(|index, component| {
934            let scale = absolute[index]
935                + relative[index] * component.abs().max(new_state.get_component(index).abs());
936            let normalized = error.get_component(index) / scale;
937            value += normalized * normalized;
938        });
939        value
940    }
941
942    fn error_norm_inf(
943        &self,
944        new_state: &Self,
945        error: &Self,
946        absolute: &BackendTolerance<f64>,
947        relative: &BackendTolerance<f64>,
948    ) -> f64 {
949        let mut value: f64 = 0.0;
950        self.for_each(|index, component| {
951            let scale = absolute[index]
952                + relative[index] * component.abs().max(new_state.get_component(index).abs());
953            value = value.max((error.get_component(index) / scale).abs());
954        });
955        value
956    }
957}
958
959struct AugmentedModel<'a, M, const N: usize, const P: usize, const W: usize> {
960    model: &'a M,
961    parameters: &'a [f64; P],
962    parameter_seeds: &'a [[f64; W]; P],
963    error: RefCell<Option<PykepError>>,
964}
965
966impl<M, const N: usize, const P: usize, const W: usize> ODE<f64, AugmentedState<N, W>>
967    for AugmentedModel<'_, M, N, P, W>
968where
969    M: DifferentiableDynamicsModel<N, P>,
970{
971    fn diff(&self, time: f64, state: &AugmentedState<N, W>, derivative: &mut AugmentedState<N, W>) {
972        if self.error.borrow().is_some() {
973            derivative.fill(f64::NAN);
974            return;
975        }
976        if let Err(error) =
977            self.model
978                .rhs(time, &state.state, self.parameters, &mut derivative.state)
979        {
980            *self.error.borrow_mut() = Some(error);
981            derivative.fill(f64::NAN);
982            return;
983        }
984        let mut state_jacobian = [[0.0; N]; N];
985        let mut parameter_jacobian = [[0.0; P]; N];
986        if let Err(error) = self.model.jacobians(
987            time,
988            &state.state,
989            self.parameters,
990            &mut state_jacobian,
991            &mut parameter_jacobian,
992        ) {
993            *self.error.borrow_mut() = Some(error);
994            derivative.fill(f64::NAN);
995            return;
996        }
997        for row in 0..N {
998            for seed in 0..W {
999                let state_term = (0..N)
1000                    .map(|column| state_jacobian[row][column] * state.sensitivities[column][seed])
1001                    .sum::<f64>();
1002                let parameter_term = (0..P)
1003                    .map(|column| {
1004                        parameter_jacobian[row][column] * self.parameter_seeds[column][seed]
1005                    })
1006                    .sum::<f64>();
1007                derivative.sensitivities[row][seed] = state_term + parameter_term;
1008            }
1009        }
1010        if derivative
1011            .state
1012            .iter()
1013            .chain(derivative.sensitivities.iter().flat_map(|row| row.iter()))
1014            .any(|value| !value.is_finite())
1015        {
1016            *self.error.borrow_mut() = Some(PykepError::IntegrationFailure {
1017                model: M::NAME,
1018                reason: "variational right-hand side returned a non-finite derivative".into(),
1019            });
1020            derivative.fill(f64::NAN);
1021        }
1022    }
1023}
1024
1025#[cfg(test)]
1026mod tests {
1027    use super::{AugmentedState, BackendState, BackendTolerance};
1028
1029    #[test]
1030    fn augmented_state_backend_layout_and_arithmetic_are_consistent() {
1031        let mut state = AugmentedState {
1032            state: [1.0, 2.0],
1033            sensitivities: [[3.0, 4.0], [5.0, 6.0]],
1034        };
1035        assert_eq!(state.len(), 6);
1036        assert_eq!(state.get_component(0), 1.0);
1037        assert_eq!(state.get_component(5), 6.0);
1038        state.set_component(0, 2.0);
1039        state.set_component(5, 7.0);
1040        assert_eq!(state.state[0], 2.0);
1041        assert_eq!(state.sensitivities[1][1], 7.0);
1042
1043        let zero = state.zeros_like();
1044        assert_eq!(zero, AugmentedState::zeros());
1045        let mut arithmetic = state.clone();
1046        arithmetic.mul_add_assign(2.0, &state);
1047        arithmetic.scale_mut(0.5);
1048        for index in 0..state.len() {
1049            assert_eq!(
1050                arithmetic.get_component(index),
1051                1.5 * state.get_component(index)
1052            );
1053        }
1054        assert!(state.norm_squared() > 0.0);
1055        assert_eq!(state.diff_norm_squared(&state), 0.0);
1056
1057        let error = AugmentedState {
1058            state: [0.1, -0.2],
1059            sensitivities: [[0.3, -0.4], [0.5, -0.6]],
1060        };
1061        let absolute = BackendTolerance::Scalar(1.0);
1062        let relative = BackendTolerance::Scalar(0.0);
1063        let squared = state.error_norm(&state, &error, &absolute, &relative);
1064        assert!((squared - 0.91).abs() < 1e-15);
1065        assert!((state.error_norm_inf(&state, &error, &absolute, &relative) - 0.6).abs() < 1e-15);
1066    }
1067}
pykep_core/integration.rs

pykep_core/
integration.rs
