Skip to main content

nyx_space/mc/
multivariate.rs

1/*
2    Nyx, blazing fast astrodynamics
3    Copyright (C) 2018-onwards Christopher Rabotin <christopher.rabotin@gmail.com>
4
5    This program is free software: you can redistribute it and/or modify
6    it under the terms of the GNU Affero General Public License as published
7    by the Free Software Foundation, either version 3 of the License, or
8    (at your option) any later version.
9
10    This program is distributed in the hope that it will be useful,
11    but WITHOUT ANY WARRANTY; without even the implied warranty of
12    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
13    GNU Affero General Public License for more details.
14
15    You should have received a copy of the GNU Affero General Public License
16    along with this program.  If not, see <https://www.gnu.org/licenses/>.
17*/
18
19use super::{DispersedState, StateDispersion};
20use crate::cosmic::AstroPhysicsSnafu;
21use crate::errors::StateError;
22use crate::md::prelude::BPlane;
23use crate::md::{AstroSnafu, StateParameter};
24use crate::{pseudo_inverse, NyxError, Spacecraft, State};
25use anise::analysis::prelude::OrbitalElement;
26use anise::astro::orbit_gradient::OrbitGrad;
27use nalgebra::{DMatrix, DVector, SMatrix, SVector};
28use rand_distr::{Distribution, Normal};
29use snafu::ResultExt;
30use std::error::Error;
31
32/// A multivariate spacecraft state generator for Monte Carlo analyses. Ensures that the covariance is properly applied on all provided state variables.
33///
34/// # Algorithm
35///
36/// The `MvnSpacecraft` allows sampling from a multivariate normal distribution defined by a template state (mean) and a set of state dispersions (uncertainties).
37/// The core difficulty is that dispersions are often defined in non-Cartesian spaces (like Keplerian orbital elements or B-Plane parameters), while the spacecraft state is represented in Cartesian coordinates (Position and Velocity).
38///
39/// The algorithm proceeds as follows:
40/// 1.  **Jacobian Computation**: It computes the Jacobian matrix `J` representing the partial derivatives of the provided dispersion parameters with respect to the Cartesian state elements (x, y, z, vx, vy, vz).
41/// 2.  **Covariance Transformation**: It constructs a diagonal covariance matrix `P` in the parameter space (assuming input dispersions are independent in that space). It then transforms this into the Cartesian covariance matrix `C` using the linear mapping approximation:
42///     `C = J_inv * P * J_inv^T`
43///     where `J_inv` is the Moore-Penrose pseudo-inverse of `J`.
44/// 3.  **Decomposition**: It performs a Singular Value Decomposition (SVD) on `C` (or the full state covariance including mass, Cr, Cd) to obtain the square root of the covariance matrix, denoted as `L`.
45///     `C = U * S * V^T = (V * sqrt(S)) * (V * sqrt(S))^T` implies `L = V * sqrt(S)`
46/// 4.  **Sampling**: To generate a sample state `X`, it draws a vector `Z` of independent standard normal variables (N(0, 1)) and applies the transformation:
47///     `X = mu + L * Z`
48///
49/// # Correctness vs. Independent Sampling
50///
51/// One might ask: "Why not just sample each Cartesian coordinate independently using a normal distribution?"
52///
53/// 1.  **Correlations**: Independent sampling of Cartesian coordinates assumes a diagonal covariance matrix, implying no correlation between position and velocity components. In orbital mechanics, states are highly correlated (e.g., velocity magnitude and radial distance are coupled by energy). `MvnSpacecraft` preserves these physical correlations by mapping the physically meaningful uncertainties (e.g., in SMA or Inclination) into the Cartesian space.
54/// 2.  **Geometry**: Uncertainties defined in orbital elements form complex shapes (like "bananas") in Cartesian space. A multivariate normal approximation in Cartesian space captures the principal axes and orientation of this uncertainty volume, which an axis-aligned bounding box (implied by independent sampling) effectively destroys.
55/// 3.  **Consistency**: By using the Jacobian transformation, we ensure that the generated samples, when mapped back to the parameter space (linearized), reproduce the input statistics (mean and standard deviation) provided by the user.
56#[derive(Clone, Debug)]
57pub struct MvnSpacecraft {
58    /// The template state
59    pub template: Spacecraft,
60    pub dispersions: Vec<StateDispersion>,
61    /// The mean of the multivariate normal distribution
62    pub mean: SVector<f64, 9>,
63    /// The dot product \sqrt{\vec s} \cdot \vec v, where S is the singular values and V the V matrix from the SVD decomp of the covariance of multivariate normal distribution
64    pub sqrt_s_v: SMatrix<f64, 9, 9>,
65    /// The standard normal distribution used to seed the multivariate normal distribution
66    pub std_norm_distr: Normal<f64>,
67}
68
69impl MvnSpacecraft {
70    /// Creates a new mulivariate state generator from a mean and covariance on the set of state parameters.
71    /// The covariance must be positive semi definite.
72    ///
73    /// # Algorithm
74    /// This function will build the rotation matrix to rotate the requested dispersions into the Spacecraft state space using [OrbitDual].
75    /// If there are any dispersions on the Cr and Cd, then these are dispersed independently (because they are iid).
76    pub fn new(
77        template: Spacecraft,
78        dispersions: Vec<StateDispersion>,
79    ) -> Result<Self, Box<dyn Error>> {
80        let mut cov = SMatrix::<f64, 9, 9>::zeros();
81        let mut mean = SVector::<f64, 9>::zeros();
82
83        let orbit_dual = OrbitGrad::from(template.orbit);
84        let mut b_plane = None;
85        for obj in &dispersions {
86            if obj.param.is_b_plane() {
87                b_plane = Some(
88                    BPlane::from_dual(orbit_dual)
89                        .context(AstroSnafu)
90                        .map_err(Box::new)?,
91                );
92                break;
93            }
94        }
95
96        let num_orbital = dispersions
97            .iter()
98            .filter(|disp| disp.param.is_orbital())
99            .count();
100
101        dbg!(num_orbital);
102
103        if num_orbital > 0 {
104            // Build the rotation matrix from the orbital dispersions to the Cartesian state.
105            let mut jac = DMatrix::from_element(num_orbital, 6, 0.0);
106            let mut covar = DMatrix::from_element(num_orbital, num_orbital, 0.0);
107            let mut means = DVector::from_element(num_orbital, 0.0);
108            let orbit_dual = OrbitGrad::from(template.orbit);
109
110            for (rno, disp) in dispersions
111                .iter()
112                .filter(|d| d.param.is_orbital())
113                .enumerate()
114            {
115                let partial = if let StateParameter::Element(oe) = disp.param {
116                    orbit_dual.partial_for(oe).context(AstroPhysicsSnafu)?
117                } else if disp.param.is_b_plane() {
118                    match disp.param {
119                        StateParameter::BdotR => b_plane.unwrap().b_r_km,
120                        StateParameter::BdotT => b_plane.unwrap().b_t_km,
121                        StateParameter::BLTOF => b_plane.unwrap().ltof_s,
122                        _ => unreachable!(),
123                    }
124                } else {
125                    unreachable!()
126                };
127
128                println!("{partial}");
129
130                for (cno, val) in [
131                    partial.wrt_x(),
132                    partial.wrt_y(),
133                    partial.wrt_z(),
134                    partial.wrt_vx(),
135                    partial.wrt_vy(),
136                    partial.wrt_vz(),
137                ]
138                .iter()
139                .copied()
140                .enumerate()
141                {
142                    jac[(rno, cno)] = val;
143                }
144                covar[(rno, rno)] = disp.std_dev.unwrap_or(0.0).powi(2);
145                means[rno] = disp.mean.unwrap_or(0.0);
146            }
147
148            // Now that we have the Jacobian that rotates from the Cartesian elements to the dispersions parameters,
149            // let's compute the inverse of this Jacobian to rotate from the dispersion params into the Cartesian elements.
150            let jac_inv = pseudo_inverse!(&jac)?;
151
152            // Rotate the orbital covariance back into the Cartesian state space, making this a 6x6.
153            let orbit_cov = &jac_inv * &covar * jac_inv.transpose();
154            // Rotate the means into the Cartesian space
155            let cartesian_mean = jac_inv * means;
156            for ii in 0..6 {
157                for jj in 0..6 {
158                    cov[(ii, jj)] = orbit_cov[(ii, jj)];
159                }
160                mean[ii] = cartesian_mean[ii];
161            }
162        };
163
164        if dispersions.len() > num_orbital {
165            for disp in &dispersions {
166                if disp.param.is_orbital() {
167                    continue;
168                } else {
169                    match disp.param {
170                        StateParameter::Cr => {
171                            cov[(7, 7)] = disp.mean.unwrap_or(0.0).powi(2);
172                        }
173                        StateParameter::Cd => {
174                            cov[(8, 8)] = disp.mean.unwrap_or(0.0).powi(2);
175                        }
176                        StateParameter::DryMass | StateParameter::PropMass => {
177                            cov[(9, 9)] = disp.mean.unwrap_or(0.0).powi(2);
178                        }
179                        _ => return Err(Box::new(StateError::ReadOnly { param: disp.param })),
180                    }
181                }
182            }
183        }
184
185        // At this point, the cov matrix is a 9x9 with all dispersions transformed into the Cartesian state space.
186        let svd = cov.svd_unordered(false, true);
187        if svd.v_t.is_none() {
188            return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
189        }
190
191        let sqrt_s = svd.singular_values.map(|x| x.sqrt());
192        let mut sqrt_s_v_t = svd.v_t.unwrap().transpose();
193
194        for (i, mut col) in sqrt_s_v_t.column_iter_mut().enumerate() {
195            col *= sqrt_s[i];
196        }
197
198        Ok(Self {
199            template,
200            dispersions,
201            mean,
202            sqrt_s_v: sqrt_s_v_t,
203            std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
204        })
205    }
206
207    /// Same as `new` but with a zero mean
208    pub fn zero_mean(
209        template: Spacecraft,
210        mut dispersions: Vec<StateDispersion>,
211    ) -> Result<Self, Box<dyn Error>> {
212        for disp in &mut dispersions {
213            disp.mean = Some(0.0);
214        }
215
216        Self::new(template, dispersions)
217    }
218
219    /// Initializes a new multivariate distribution using the state data in the spacecraft state space.
220    pub fn from_spacecraft_cov(
221        template: Spacecraft,
222        cov: SMatrix<f64, 9, 9>,
223        mean: SVector<f64, 9>,
224    ) -> Result<Self, Box<dyn Error>> {
225        // Check that covariance is PSD by ensuring that all the eigenvalues are positive or nil
226        match cov.eigenvalues() {
227            None => return Err(Box::new(NyxError::CovarianceMatrixNotPsd)),
228            Some(evals) => {
229                for eigenval in &evals {
230                    if *eigenval < 0.0 {
231                        return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
232                    }
233                }
234            }
235        };
236
237        let svd = cov.svd_unordered(false, true);
238        if svd.v_t.is_none() {
239            return Err(Box::new(NyxError::CovarianceMatrixNotPsd));
240        }
241
242        let s = svd.singular_values;
243        // Item by item multiplication
244        let mut sqrt_s_v = svd.v_t.unwrap().transpose();
245        for (i, mut col) in sqrt_s_v.column_iter_mut().enumerate() {
246            col *= s[i].sqrt();
247        }
248
249        let dispersions = vec![
250            StateDispersion::builder()
251                .param(StateParameter::Element(OrbitalElement::X))
252                .std_dev(cov[(0, 0)])
253                .build(),
254            StateDispersion::builder()
255                .param(StateParameter::Element(OrbitalElement::Y))
256                .std_dev(cov[(1, 1)])
257                .build(),
258            StateDispersion::builder()
259                .param(StateParameter::Element(OrbitalElement::Z))
260                .std_dev(cov[(2, 2)])
261                .build(),
262            StateDispersion::builder()
263                .param(StateParameter::Element(OrbitalElement::VX))
264                .std_dev(cov[(3, 3)])
265                .build(),
266            StateDispersion::builder()
267                .param(StateParameter::Element(OrbitalElement::VY))
268                .std_dev(cov[(4, 4)])
269                .build(),
270            StateDispersion::builder()
271                .param(StateParameter::Element(OrbitalElement::VZ))
272                .std_dev(cov[(5, 5)])
273                .build(),
274            StateDispersion::builder()
275                .param(StateParameter::Cr)
276                .std_dev(cov[(6, 6)])
277                .build(),
278            StateDispersion::builder()
279                .param(StateParameter::Cd)
280                .std_dev(cov[(7, 7)])
281                .build(),
282            StateDispersion::builder()
283                .param(StateParameter::PropMass)
284                .std_dev(cov[(8, 8)])
285                .build(),
286        ];
287
288        Ok(Self {
289            template,
290            dispersions,
291            mean,
292            sqrt_s_v,
293            std_norm_distr: Normal::new(0.0, 1.0).unwrap(),
294        })
295    }
296}
297
298impl Distribution<DispersedState<Spacecraft>> for MvnSpacecraft {
299    fn sample<R: rand::Rng + ?Sized>(&self, rng: &mut R) -> DispersedState<Spacecraft> {
300        // Generate the vector representing the state
301        let x_rng = SVector::<f64, 9>::from_fn(|_, _| self.std_norm_distr.sample(rng));
302        let x = self.sqrt_s_v * x_rng + self.mean;
303        let mut state = self.template;
304
305        // Set the new state data
306        for (coord, val) in x.iter().copied().enumerate() {
307            if coord < 3 {
308                state.orbit.radius_km[coord] += val;
309            } else if coord < 6 {
310                state.orbit.velocity_km_s[coord % 3] += val;
311            } else if coord == 6 {
312                state.srp.coeff_reflectivity += val;
313            } else if coord == 7 {
314                state.drag.coeff_drag += val;
315            } else if coord == 8 {
316                state.mass.prop_mass_kg += val;
317            }
318        }
319
320        let mut actual_dispersions = Vec::new();
321        for disp in &self.dispersions {
322            // Compute the delta
323            let delta = self.template.value(disp.param).unwrap() - state.value(disp.param).unwrap();
324            actual_dispersions.push((disp.param, delta));
325        }
326
327        DispersedState {
328            state,
329            actual_dispersions,
330        }
331    }
332}
333
334#[cfg(test)]
335mod multivariate_ut {
336    use super::*;
337    use crate::time::Epoch;
338    use crate::Spacecraft;
339    use crate::GMAT_EARTH_GM;
340    use anise::constants::frames::EARTH_J2000;
341    use anise::prelude::Orbit;
342    use rand::RngExt;
343    use statrs;
344    use statrs::distribution::ContinuousCDF;
345
346    /// Returns the Mahalanobis distance of a random variable x from a multivariate normal distribution
347    /// with mean mu and covariance matrix S.
348    fn mahalanobis_distance<const T: usize>(
349        x: &SVector<f64, T>,
350        mu: &SVector<f64, T>,
351        cov_inv: &SMatrix<f64, T, T>,
352    ) -> f64 {
353        let delta = x - mu;
354        (delta.transpose() * cov_inv * delta)[(0, 0)]
355    }
356
357    #[test]
358    fn mahalanobis_distance_test() {
359        // Simple 2D case with diagonal covariance
360        let x = SVector::<f64, 2>::new(1.0, 2.0);
361        let mu = SVector::<f64, 2>::new(0.0, 0.0);
362        let cov = SMatrix::<f64, 2, 2>::from_diagonal(&SVector::<f64, 2>::new(2.0, 3.0));
363        let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
364
365        // Expected distance: (1-0)^2/2 + (2-0)^2/3 = 1/2 + 4/3 = 0.5 + 1.333... = 1.8333...
366        let md = mahalanobis_distance(&x, &mu, &cov_inv);
367        assert!((md - 1.8333333333333333).abs() < 1e-9);
368    }
369
370    #[test]
371    fn test_mvn_generator() {
372        // Generate a random covariance matrix.
373        let mut rng = rand::rng();
374        let cov = SMatrix::<f64, 6, 6>::from_fn(|_, _| rng.random());
375        let cov = cov * cov.transpose();
376        let mut cov_resized = SMatrix::<f64, 9, 9>::zeros();
377        cov_resized.fixed_view_mut::<6, 6>(0, 0).copy_from(&cov);
378
379        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
380
381        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
382        let state = Spacecraft::builder()
383            .orbit(Orbit::keplerian(
384                8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
385            ))
386            .build();
387
388        let mvn = MvnSpacecraft::from_spacecraft_cov(state, cov_resized, SVector::zeros()).unwrap();
389
390        // Generate N samples from this distribution
391        let n = 1000;
392        let samples = mvn.sample_iter(&mut rng).take(n);
393
394        // Compute the Mahalanobis distance for each sample
395        let cov_inv = cov_resized.pseudo_inverse(1e-12).unwrap();
396        let md: Vec<f64> = samples
397            .map(|sample| {
398                let mut v = SVector::<f64, 9>::zeros();
399                for i in 0..6 {
400                    v[i] = sample.state.orbit.to_cartesian_pos_vel()[i]
401                        - state.orbit.to_cartesian_pos_vel()[i];
402                }
403                mahalanobis_distance(&v, &SVector::zeros(), &cov_inv)
404            })
405            .collect();
406
407        // Perform a chi-squared test on the Mahalanobis distances.
408        // The squared Mahalanobis distance of a sample from the mean of a multivariate normal distribution
409        // with k degrees of freedom follows a chi-squared distribution with k degrees of freedom.
410        // We will test if the 95th percentile of the Mahalanobis distances is within the 95th percentile
411        // of the chi-squared distribution.
412        let mut md = md;
413        md.sort_by(|a, b| a.partial_cmp(b).unwrap());
414        let p95_md = md[(0.95 * n as f64) as usize];
415
416        let chi_squared = statrs::distribution::ChiSquared::new(6.0).unwrap();
417        let p95_chi_squared = chi_squared.inverse_cdf(0.95);
418
419        assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
420    }
421
422    #[test]
423    fn disperse_r_mag() {
424        use anise::constants::frames::EARTH_J2000;
425        use anise::prelude::Orbit;
426
427        use crate::time::Epoch;
428        use rand_pcg::Pcg64Mcg;
429
430        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
431        // Create a reproducible fast seed
432        let seed = 0;
433
434        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
435
436        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
437        let state = Spacecraft::builder()
438            .orbit(Orbit::keplerian(
439                8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k,
440            ))
441            .build();
442
443        // Check that we can modify the radius magnitude
444        let std_dev = 1.0;
445        let generator = MvnSpacecraft::new(
446            state,
447            vec![StateDispersion::builder()
448                .param(StateParameter::Element(OrbitalElement::Rmag))
449                .std_dev(std_dev)
450                .build()],
451        )
452        .unwrap();
453
454        let rng = Pcg64Mcg::new(seed);
455        let init_rmag = state.orbit.rmag_km();
456        let cnt_too_far: u16 = generator
457            .sample_iter(rng)
458            .take(1000)
459            .map(|dispersed_state| {
460                if (init_rmag - dispersed_state.state.orbit.rmag_km()).abs() >= 3.0 * std_dev {
461                    1
462                } else {
463                    0
464                }
465            })
466            .sum::<u16>();
467
468        // We specified a seed so we know exactly what to expect and we've reset the seed to 0.
469        assert_eq!(
470            cnt_too_far,
471            6, // Mathematically, this should be 3!
472            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
473        );
474    }
475
476    #[test]
477    fn disperse_full_cartesian() {
478        use anise::constants::frames::EARTH_J2000;
479        use anise::prelude::Orbit;
480
481        use crate::time::Epoch;
482        use crate::Spacecraft;
483        use crate::GMAT_EARTH_GM;
484
485        use rand_pcg::Pcg64Mcg;
486
487        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
488
489        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
490        let state = Orbit::keplerian(8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k);
491
492        let std_dev = [10.0, 10.0, 10.0, 0.2, 0.2, 0.2, 0.0, 0.0, 0.0];
493
494        let generator = MvnSpacecraft::new(
495            Spacecraft {
496                orbit: state,
497                ..Default::default()
498            },
499            vec![
500                StateDispersion::builder()
501                    .param(StateParameter::Element(OrbitalElement::X))
502                    .std_dev(10.0)
503                    .build(),
504                StateDispersion::builder()
505                    .param(StateParameter::Element(OrbitalElement::Y))
506                    .std_dev(10.0)
507                    .build(),
508                StateDispersion::builder()
509                    .param(StateParameter::Element(OrbitalElement::Z))
510                    .std_dev(10.0)
511                    .build(),
512                StateDispersion::builder()
513                    .param(StateParameter::Element(OrbitalElement::VX))
514                    .std_dev(0.2)
515                    .build(),
516                StateDispersion::builder()
517                    .param(StateParameter::Element(OrbitalElement::VY))
518                    .std_dev(0.2)
519                    .build(),
520                StateDispersion::builder()
521                    .param(StateParameter::Element(OrbitalElement::VZ))
522                    .std_dev(0.2)
523                    .build(),
524            ],
525        )
526        .unwrap();
527
528        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
529        // Create a reproducible fast seed
530        let seed = 0;
531        let rng = Pcg64Mcg::new(seed);
532
533        let cnt_too_far: u16 = generator
534            .sample_iter(rng)
535            .take(1000)
536            .map(|dispersed_state| {
537                let mut cnt = 0;
538                for (idx, val_std_dev) in std_dev.iter().take(6).enumerate() {
539                    let cur_val = dispersed_state.state.to_vector()[idx];
540                    let nom_val = state.to_cartesian_pos_vel()[idx];
541                    if (cur_val - nom_val).abs() > *val_std_dev {
542                        cnt += 1;
543                    }
544                }
545                cnt
546            })
547            .sum::<u16>();
548
549        // We specified a seed so we know exactly what to expect
550        assert_eq!(
551            cnt_too_far / 6,
552            312,
553            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
554        );
555    }
556
557    #[test]
558    fn disperse_raan_only() {
559        use anise::constants::frames::EARTH_J2000;
560        use anise::prelude::Orbit;
561
562        use crate::time::Epoch;
563        use rand_pcg::Pcg64Mcg;
564
565        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
566
567        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
568        let state = Spacecraft::builder()
569            .orbit(Orbit::keplerian(
570                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
571            ))
572            .build();
573
574        let angle_sigma_deg = 0.2;
575
576        assert!(StateParameter::Element(OrbitalElement::RAAN).is_orbital());
577
578        let generator = MvnSpacecraft::new(
579            state,
580            vec![StateDispersion::zero_mean(
581                StateParameter::Element(OrbitalElement::RAAN),
582                angle_sigma_deg,
583            )],
584        )
585        .unwrap();
586
587        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
588        // Create a reproducible fast seed
589        let seed = 0;
590        let rng = Pcg64Mcg::new(seed);
591        let n = 1000;
592        let samples = generator.sample_iter(rng).take(n);
593
594        let cov = SMatrix::<f64, 1, 1>::from_diagonal(&SVector::<f64, 1>::from_vec(vec![
595            angle_sigma_deg.powi(2),
596        ]));
597        let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
598
599        let md: Vec<f64> = samples
600            .map(|dispersed_state| {
601                // For all other orbital parameters, make sure that we have not changed things dramatically.
602                for param in [
603                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
604                    StateParameter::Element(OrbitalElement::Inclination),
605                ] {
606                    let orig = state.value(param).unwrap();
607                    let new = dispersed_state.state.value(param).unwrap();
608                    let prct_change = 100.0 * (orig - new).abs() / orig;
609                    assert!(
610                        prct_change < 5.0,
611                        "{param} changed by {prct_change:.3} % ({orig:.3e} -> {new:.3e})"
612                    );
613                }
614
615                let delta =
616                    SVector::<f64, 1>::from_vec(vec![dispersed_state.actual_dispersions[0].1]);
617                mahalanobis_distance(&delta, &SVector::zeros(), &cov_inv)
618            })
619            .collect();
620
621        let mut md = md;
622        md.sort_by(|a, b| a.partial_cmp(b).unwrap());
623        let p95_md = md[(0.95 * n as f64) as usize];
624
625        let chi_squared = statrs::distribution::ChiSquared::new(1.0).unwrap();
626        let p95_chi_squared = chi_squared.inverse_cdf(0.95);
627
628        assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
629    }
630
631    #[test]
632    fn disperse_keplerian() {
633        use anise::constants::frames::EARTH_J2000;
634        use anise::prelude::Orbit;
635
636        use crate::time::Epoch;
637        use rand_pcg::Pcg64Mcg;
638
639        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
640
641        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
642        let state = Spacecraft::builder()
643            .orbit(Orbit::keplerian(
644                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
645            ))
646            .build();
647
648        let sma_sigma_km = 10.0;
649        let inc_sigma_deg = 0.15;
650        let angle_sigma_deg = 0.02;
651
652        let generator = MvnSpacecraft::new(
653            state,
654            vec![
655                StateDispersion::zero_mean(
656                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
657                    sma_sigma_km,
658                ),
659                StateDispersion::zero_mean(
660                    StateParameter::Element(OrbitalElement::Inclination),
661                    inc_sigma_deg,
662                ),
663                StateDispersion::zero_mean(
664                    StateParameter::Element(OrbitalElement::RAAN),
665                    angle_sigma_deg,
666                ),
667                StateDispersion::zero_mean(
668                    StateParameter::Element(OrbitalElement::AoP),
669                    angle_sigma_deg,
670                ),
671            ],
672        )
673        .unwrap();
674
675        // The generator computes the equivalent Cartesian covariance. Let's retrieve it.
676        let expected_cart_cov = &generator.sqrt_s_v * generator.sqrt_s_v.transpose();
677
678        // Create a reproducible fast seed, and generate the samples in the Cartesian space
679        let seed = 0;
680        let rng = Pcg64Mcg::new(seed);
681        let n = 2000; // Increase samples for better stats
682        let samples: Vec<SVector<f64, 6>> = generator
683            .sample_iter(rng)
684            .take(n)
685            .map(|s| s.state.orbit.to_cartesian_pos_vel())
686            .collect();
687
688        let nominal_cart = state.orbit.to_cartesian_pos_vel();
689
690        // Check that the mean of the Cartesian states is close to the nominal
691        let mut mean_cart = SVector::<f64, 6>::zeros();
692        for sample in &samples {
693            mean_cart += sample;
694        }
695        mean_cart /= n as f64;
696
697        let mean_diff = mean_cart - nominal_cart;
698        // We expect some deviation because of the non-linearities.
699        assert!(mean_diff.norm() < 1.0);
700
701        // Check that the covariance of the Cartesian states is close to the expected one
702        let mut sample_cov = SMatrix::<f64, 6, 6>::zeros();
703        for sample in &samples {
704            let disp_vec = sample - &mean_cart;
705            sample_cov += &disp_vec * disp_vec.transpose();
706        }
707        sample_cov /= (n - 1) as f64;
708
709        let expected_cov_6x6 = expected_cart_cov.fixed_view::<6, 6>(0, 0);
710
711        let diff = sample_cov - expected_cov_6x6;
712        assert!(diff.norm() / expected_cov_6x6.norm() < 0.2);
713    }
714}