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::{NyxError, Spacecraft, State, pseudo_inverse};
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#[cfg(feature = "python")]
33use pyo3::prelude::*;
34
35/// A multivariate spacecraft state generator for Monte Carlo analyses. Ensures that the covariance is properly applied on all provided state variables.
36///
37/// # Algorithm
38///
39/// The `MvnSpacecraft` allows sampling from a multivariate normal distribution defined by a template state (mean) and a set of state dispersions (uncertainties).
40/// 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).
41///
42/// The algorithm proceeds as follows:
43/// 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).
44/// 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:
45///     `C = J_inv * P * J_inv^T`
46///     where `J_inv` is the Moore-Penrose pseudo-inverse of `J`.
47/// 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`.
48///     `C = U * S * V^T = (V * sqrt(S)) * (V * sqrt(S))^T` implies `L = V * sqrt(S)`
49/// 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:
50///     `X = mu + L * Z`
51///
52/// # Correctness vs. Independent Sampling
53///
54/// One might ask: "Why not just sample each Cartesian coordinate independently using a normal distribution?"
55///
56/// 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.
57/// 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.
58/// 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.
59#[derive(Clone, Debug)]
60#[cfg_attr(feature = "python", pyclass(from_py_object))]
61pub struct MvnSpacecraft {
62    /// The template state
63    pub template: Spacecraft,
64    pub dispersions: Vec<StateDispersion>,
65    /// The mean of the multivariate normal distribution
66    pub mean: SVector<f64, 9>,
67    /// 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
68    pub sqrt_s_v: SMatrix<f64, 9, 9>,
69    /// The standard normal distribution used to seed the multivariate normal distribution
70    pub std_norm_distr: Normal<f64>,
71}
72
73impl MvnSpacecraft {
74    /// Creates a new mulivariate state generator from a mean and covariance on the set of state parameters.
75    /// The covariance must be positive semi definite.
76    ///
77    /// # Algorithm
78    /// This function will build the rotation matrix to rotate the requested dispersions into the Spacecraft state space using [OrbitDual].
79    /// If there are any dispersions on the Cr and Cd, then these are dispersed independently (because they are iid).
80    pub fn new(
81        template: Spacecraft,
82        dispersions: Vec<StateDispersion>,
83    ) -> Result<Self, Box<dyn Error>> {
84        let mut cov = SMatrix::<f64, 9, 9>::zeros();
85        let mut mean = SVector::<f64, 9>::zeros();
86
87        let orbit_dual = OrbitGrad::from(template.orbit);
88        let mut b_plane = None;
89        for obj in &dispersions {
90            if obj.param.is_b_plane() {
91                b_plane = Some(
92                    BPlane::from_dual(orbit_dual)
93                        .context(AstroSnafu)
94                        .map_err(Box::new)?,
95                );
96                break;
97            }
98        }
99
100        let num_orbital = dispersions
101            .iter()
102            .filter(|disp| disp.param.is_orbital())
103            .count();
104
105        if num_orbital > 0 {
106            // Build the rotation matrix from the orbital dispersions to the Cartesian state.
107            let mut jac = DMatrix::from_element(num_orbital, 6, 0.0);
108            let mut covar = DMatrix::from_element(num_orbital, num_orbital, 0.0);
109            let mut means = DVector::from_element(num_orbital, 0.0);
110            let orbit_dual = OrbitGrad::from(template.orbit);
111
112            for (rno, disp) in dispersions
113                .iter()
114                .filter(|d| d.param.is_orbital())
115                .enumerate()
116            {
117                let partial = if let StateParameter::Element(oe) = disp.param {
118                    orbit_dual.partial_for(oe).context(AstroPhysicsSnafu)?
119                } else if disp.param.is_b_plane() {
120                    match disp.param {
121                        StateParameter::BdotR() => b_plane.unwrap().b_r_km,
122                        StateParameter::BdotT() => b_plane.unwrap().b_t_km,
123                        StateParameter::BLTOF() => b_plane.unwrap().ltof_s,
124                        _ => unreachable!(),
125                    }
126                } else {
127                    unreachable!()
128                };
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::GMAT_EARTH_GM;
338    use crate::Spacecraft;
339    use crate::time::Epoch;
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![
448                StateDispersion::builder()
449                    .param(StateParameter::Element(OrbitalElement::Rmag))
450                    .std_dev(std_dev)
451                    .build(),
452            ],
453        )
454        .unwrap();
455
456        let rng = Pcg64Mcg::new(seed);
457        let init_rmag = state.orbit.rmag_km();
458        let cnt_too_far: u16 = generator
459            .sample_iter(rng)
460            .take(1000)
461            .map(|dispersed_state| {
462                if (init_rmag - dispersed_state.state.orbit.rmag_km()).abs() >= 3.0 * std_dev {
463                    1
464                } else {
465                    0
466                }
467            })
468            .sum::<u16>();
469
470        // We specified a seed so we know exactly what to expect and we've reset the seed to 0.
471        assert_eq!(
472            cnt_too_far,
473            6, // Mathematically, this should be 3!
474            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
475        );
476    }
477
478    #[test]
479    fn disperse_full_cartesian() {
480        use anise::constants::frames::EARTH_J2000;
481        use anise::prelude::Orbit;
482
483        use crate::GMAT_EARTH_GM;
484        use crate::Spacecraft;
485        use crate::time::Epoch;
486
487        use rand_pcg::Pcg64Mcg;
488
489        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
490
491        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
492        let state = Orbit::keplerian(8_191.93, 1e-6, 12.85, 306.614, 314.19, 99.887_7, dt, eme2k);
493
494        let std_dev = [10.0, 10.0, 10.0, 0.2, 0.2, 0.2, 0.0, 0.0, 0.0];
495
496        let generator = MvnSpacecraft::new(
497            Spacecraft {
498                orbit: state,
499                ..Default::default()
500            },
501            vec![
502                StateDispersion::builder()
503                    .param(StateParameter::Element(OrbitalElement::X))
504                    .std_dev(10.0)
505                    .build(),
506                StateDispersion::builder()
507                    .param(StateParameter::Element(OrbitalElement::Y))
508                    .std_dev(10.0)
509                    .build(),
510                StateDispersion::builder()
511                    .param(StateParameter::Element(OrbitalElement::Z))
512                    .std_dev(10.0)
513                    .build(),
514                StateDispersion::builder()
515                    .param(StateParameter::Element(OrbitalElement::VX))
516                    .std_dev(0.2)
517                    .build(),
518                StateDispersion::builder()
519                    .param(StateParameter::Element(OrbitalElement::VY))
520                    .std_dev(0.2)
521                    .build(),
522                StateDispersion::builder()
523                    .param(StateParameter::Element(OrbitalElement::VZ))
524                    .std_dev(0.2)
525                    .build(),
526            ],
527        )
528        .unwrap();
529
530        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
531        // Create a reproducible fast seed
532        let seed = 0;
533        let rng = Pcg64Mcg::new(seed);
534
535        let cnt_too_far: u16 = generator
536            .sample_iter(rng)
537            .take(1000)
538            .map(|dispersed_state| {
539                let mut cnt = 0;
540                for (idx, val_std_dev) in std_dev.iter().take(6).enumerate() {
541                    let cur_val = dispersed_state.state.to_vector()[idx];
542                    let nom_val = state.to_cartesian_pos_vel()[idx];
543                    if (cur_val - nom_val).abs() > *val_std_dev {
544                        cnt += 1;
545                    }
546                }
547                cnt
548            })
549            .sum::<u16>();
550
551        // We specified a seed so we know exactly what to expect
552        assert_eq!(
553            cnt_too_far / 6,
554            312,
555            "Should have about 3% of samples being more than 3 sigma away, got {cnt_too_far}"
556        );
557    }
558
559    #[test]
560    fn disperse_raan_only() {
561        use anise::constants::frames::EARTH_J2000;
562        use anise::prelude::Orbit;
563
564        use crate::time::Epoch;
565        use rand_pcg::Pcg64Mcg;
566
567        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
568
569        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
570        let state = Spacecraft::builder()
571            .orbit(Orbit::keplerian(
572                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
573            ))
574            .build();
575
576        let angle_sigma_deg = 0.2;
577
578        assert!(StateParameter::Element(OrbitalElement::RAAN).is_orbital());
579
580        let generator = MvnSpacecraft::new(
581            state,
582            vec![StateDispersion::zero_mean(
583                StateParameter::Element(OrbitalElement::RAAN),
584                angle_sigma_deg,
585            )],
586        )
587        .unwrap();
588
589        // Ensure that this worked: a 3 sigma deviation around 1 km means we shouldn't have 99.7% of samples within those bounds.
590        // Create a reproducible fast seed
591        let seed = 0;
592        let rng = Pcg64Mcg::new(seed);
593        let n = 1000;
594        let samples = generator.sample_iter(rng).take(n);
595
596        let cov = SMatrix::<f64, 1, 1>::from_diagonal(&SVector::<f64, 1>::from_vec(vec![
597            angle_sigma_deg.powi(2),
598        ]));
599        let cov_inv = cov.pseudo_inverse(1e-12).unwrap();
600
601        let md: Vec<f64> = samples
602            .map(|dispersed_state| {
603                // For all other orbital parameters, make sure that we have not changed things dramatically.
604                for param in [
605                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
606                    StateParameter::Element(OrbitalElement::Inclination),
607                ] {
608                    let orig = state.value(param).unwrap();
609                    let new = dispersed_state.state.value(param).unwrap();
610                    let prct_change = 100.0 * (orig - new).abs() / orig;
611                    assert!(
612                        prct_change < 5.0,
613                        "{param} changed by {prct_change:.3} % ({orig:.3e} -> {new:.3e})"
614                    );
615                }
616
617                let delta =
618                    SVector::<f64, 1>::from_vec(vec![dispersed_state.actual_dispersions[0].1]);
619                mahalanobis_distance(&delta, &SVector::zeros(), &cov_inv)
620            })
621            .collect();
622
623        let mut md = md;
624        md.sort_by(|a, b| a.partial_cmp(b).unwrap());
625        let p95_md = md[(0.95 * n as f64) as usize];
626
627        let chi_squared = statrs::distribution::ChiSquared::new(1.0).unwrap();
628        let p95_chi_squared = chi_squared.inverse_cdf(0.95);
629
630        assert!((p95_md - p95_chi_squared).abs() / p95_chi_squared < 0.2);
631    }
632
633    #[test]
634    fn disperse_keplerian() {
635        use anise::constants::frames::EARTH_J2000;
636        use anise::prelude::Orbit;
637
638        use crate::time::Epoch;
639        use rand_pcg::Pcg64Mcg;
640
641        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
642
643        let dt = Epoch::from_gregorian_utc_at_midnight(2021, 1, 31);
644        let state = Spacecraft::builder()
645            .orbit(Orbit::keplerian(
646                8_100.0, 1e-6, 12.85, 356.614, 14.19, 199.887_7, dt, eme2k,
647            ))
648            .build();
649
650        let sma_sigma_km = 10.0;
651        let inc_sigma_deg = 0.15;
652        let angle_sigma_deg = 0.02;
653
654        let generator = MvnSpacecraft::new(
655            state,
656            vec![
657                StateDispersion::zero_mean(
658                    StateParameter::Element(OrbitalElement::SemiMajorAxis),
659                    sma_sigma_km,
660                ),
661                StateDispersion::zero_mean(
662                    StateParameter::Element(OrbitalElement::Inclination),
663                    inc_sigma_deg,
664                ),
665                StateDispersion::zero_mean(
666                    StateParameter::Element(OrbitalElement::RAAN),
667                    angle_sigma_deg,
668                ),
669                StateDispersion::zero_mean(
670                    StateParameter::Element(OrbitalElement::AoP),
671                    angle_sigma_deg,
672                ),
673            ],
674        )
675        .unwrap();
676
677        // The generator computes the equivalent Cartesian covariance. Let's retrieve it.
678        let expected_cart_cov = generator.sqrt_s_v * generator.sqrt_s_v.transpose();
679
680        // Create a reproducible fast seed, and generate the samples in the Cartesian space
681        let seed = 0;
682        let rng = Pcg64Mcg::new(seed);
683        let n = 2000; // Increase samples for better stats
684        let samples: Vec<SVector<f64, 6>> = generator
685            .sample_iter(rng)
686            .take(n)
687            .map(|s| s.state.orbit.to_cartesian_pos_vel())
688            .collect();
689
690        let nominal_cart = state.orbit.to_cartesian_pos_vel();
691
692        // Check that the mean of the Cartesian states is close to the nominal
693        let mut mean_cart = SVector::<f64, 6>::zeros();
694        for sample in &samples {
695            mean_cart += sample;
696        }
697        mean_cart /= n as f64;
698
699        let mean_diff = mean_cart - nominal_cart;
700        // We expect some deviation because of the non-linearities.
701        assert!(mean_diff.norm() < 1.0);
702
703        // Check that the covariance of the Cartesian states is close to the expected one
704        let mut sample_cov = SMatrix::<f64, 6, 6>::zeros();
705        for sample in &samples {
706            let disp_vec = sample - mean_cart;
707            sample_cov += disp_vec * disp_vec.transpose();
708        }
709        sample_cov /= (n - 1) as f64;
710
711        let expected_cov_6x6 = expected_cart_cov.fixed_view::<6, 6>(0, 0);
712
713        let diff = sample_cov - expected_cov_6x6;
714        assert!(diff.norm() / expected_cov_6x6.norm() < 0.2);
715    }
716}