1use 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#[derive(Clone, Debug)]
57pub struct MvnSpacecraft {
58 pub template: Spacecraft,
60 pub dispersions: Vec<StateDispersion>,
61 pub mean: SVector<f64, 9>,
63 pub sqrt_s_v: SMatrix<f64, 9, 9>,
65 pub std_norm_distr: Normal<f64>,
67}
68
69impl MvnSpacecraft {
70 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 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 let jac_inv = pseudo_inverse!(&jac)?;
151
152 let orbit_cov = &jac_inv * &covar * jac_inv.transpose();
154 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 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 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 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 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 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 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 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 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 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 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 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 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 let n = 1000;
392 let samples = mvn.sample_iter(&mut rng).take(n);
393
394 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 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 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 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 assert_eq!(
470 cnt_too_far,
471 6, "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 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 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 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 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 let expected_cart_cov = &generator.sqrt_s_v * generator.sqrt_s_v.transpose();
677
678 let seed = 0;
680 let rng = Pcg64Mcg::new(seed);
681 let n = 2000; 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 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 assert!(mean_diff.norm() < 1.0);
700
701 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}