nyx_space/od/estimate/
kfestimate.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::{Estimate, State};
20use crate::cosmic::{AstroError, AstroPhysicsSnafu};
21use crate::linalg::allocator::Allocator;
22use crate::linalg::{DefaultAllocator, DimName, Matrix, OMatrix, OVector};
23use crate::mc::{MvnSpacecraft, StateDispersion};
24use crate::Spacecraft;
25use anise::analysis::prelude::OrbitalElement;
26use anise::astro::orbit_gradient::OrbitGrad;
27use nalgebra::Const;
28use nalgebra::SMatrix;
29use rand::rngs::SysRng;
30use rand::SeedableRng;
31use rand_distr::Distribution;
32use rand_pcg::Pcg64Mcg;
33use snafu::ResultExt;
34use std::cmp::PartialEq;
35use std::error::Error;
36use std::fmt;
37use std::ops::Mul;
38
39/// Kalman filter Estimate
40#[derive(Debug, Copy, Clone, PartialEq)]
41pub struct KfEstimate<T: State>
42where
43    DefaultAllocator: Allocator<<T as State>::Size>
44        + Allocator<<T as State>::Size, <T as State>::Size>
45        + Allocator<<T as State>::Size>
46        + Allocator<<T as State>::VecLength>,
47    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
48    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
49{
50    /// The estimated state
51    pub nominal_state: T,
52    /// The state deviation
53    pub state_deviation: OVector<f64, <T as State>::Size>,
54    /// The Covariance of this estimate
55    pub covar: OMatrix<f64, <T as State>::Size, <T as State>::Size>,
56    /// The predicted covariance of this estimate
57    pub covar_bar: OMatrix<f64, <T as State>::Size, <T as State>::Size>,
58    /// Whether or not this is a predicted estimate from a time update, or an estimate from a measurement
59    pub predicted: bool,
60    /// The STM used to compute this Estimate
61    pub stm: OMatrix<f64, <T as State>::Size, <T as State>::Size>,
62}
63
64impl<T: State> KfEstimate<T>
65where
66    DefaultAllocator: Allocator<<T as State>::Size>
67        + Allocator<<T as State>::Size, <T as State>::Size>
68        + Allocator<<T as State>::Size>
69        + Allocator<<T as State>::VecLength>,
70    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
71    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
72{
73    /// Initializes a new filter estimate from the nominal state (not dispersed) and the full covariance
74    pub fn from_covar(
75        nominal_state: T,
76        covar: OMatrix<f64, <T as State>::Size, <T as State>::Size>,
77    ) -> Self {
78        Self {
79            nominal_state,
80            state_deviation: OVector::<f64, <T as State>::Size>::zeros(),
81            covar,
82            covar_bar: covar,
83            predicted: true,
84            stm: OMatrix::<f64, <T as State>::Size, <T as State>::Size>::identity(),
85        }
86    }
87
88    /// Initializes a new filter estimate from the nominal state (not dispersed) and the diagonal of the covariance
89    pub fn from_diag(nominal_state: T, diag: OVector<f64, <T as State>::Size>) -> Self {
90        let covar = Matrix::from_diagonal(&diag);
91        Self {
92            nominal_state,
93            state_deviation: OVector::<f64, <T as State>::Size>::zeros(),
94            covar,
95            covar_bar: covar,
96            predicted: true,
97            stm: OMatrix::<f64, <T as State>::Size, <T as State>::Size>::identity(),
98        }
99    }
100}
101
102impl KfEstimate<Spacecraft> {
103    /// Generates an initial Kalman filter state estimate dispersed from the nominal state using the provided standard deviation parameters.
104    ///
105    /// The resulting estimate will have a diagonal covariance matrix constructed from the variances of each parameter.
106    /// *Limitation:* This method may not work correctly for all Keplerian orbital elements, refer to
107    /// <https://github.com/nyx-space/nyx/issues/339> for details.
108    pub fn disperse_from_diag(
109        nominal_state: Spacecraft,
110        dispersions: Vec<StateDispersion>,
111        seed: Option<u128>,
112    ) -> Result<Self, Box<dyn Error>> {
113        let generator = MvnSpacecraft::new(nominal_state, dispersions)?;
114
115        let mut rng = match seed {
116            Some(seed) => Pcg64Mcg::new(seed),
117            None => Pcg64Mcg::try_from_rng(&mut SysRng).unwrap(),
118        };
119        let dispersed_state = generator.sample(&mut rng);
120
121        // Compute the difference between both states
122        let delta_orbit = (nominal_state.orbit - dispersed_state.state.orbit).unwrap();
123
124        // Build the covariance as three times the absolute value of the error, squared.
125
126        let diag_data = [
127            (3.0 * delta_orbit.radius_km.x.abs()).powi(2),
128            (3.0 * delta_orbit.radius_km.y.abs()).powi(2),
129            (3.0 * delta_orbit.radius_km.z.abs()).powi(2),
130            (3.0 * delta_orbit.velocity_km_s.x.abs()).powi(2),
131            (3.0 * delta_orbit.velocity_km_s.y.abs()).powi(2),
132            (3.0 * delta_orbit.velocity_km_s.z.abs()).powi(2),
133            (3.0 * (nominal_state.srp.coeff_reflectivity
134                - dispersed_state.state.srp.coeff_reflectivity)
135                .abs())
136            .powi(2),
137            (3.0 * (nominal_state.drag.coeff_drag - dispersed_state.state.drag.coeff_drag).abs())
138                .powi(2),
139            (3.0 * (nominal_state.mass.prop_mass_kg - dispersed_state.state.mass.prop_mass_kg)
140                .abs())
141            .powi(2),
142        ];
143
144        let diag = OVector::<f64, Const<9>>::from_iterator(diag_data);
145
146        // Build the covar from the diagonal
147        let covar = Matrix::from_diagonal(&diag);
148
149        Ok(Self {
150            nominal_state: dispersed_state.state,
151            state_deviation: OVector::<f64, Const<9>>::zeros(),
152            covar,
153            covar_bar: covar,
154            predicted: true,
155            stm: OMatrix::<f64, Const<9>, Const<9>>::identity(),
156        })
157    }
158
159    /// Builds a multivariate random variable from this estimate's nominal state and covariance, zero mean.
160    pub fn to_random_variable(&self) -> Result<MvnSpacecraft, Box<dyn Error>> {
161        MvnSpacecraft::from_spacecraft_cov(self.nominal_state, self.covar, self.state_deviation)
162    }
163
164    /// Returns the 1-sigma uncertainty for a given parameter, in that parameter's unit
165    ///
166    /// This method uses the [OrbitDual] structure to compute the estimate in the hyperdual space
167    /// and rotate the nominal covariance into that space.
168    pub fn sigma_for(&self, param: OrbitalElement) -> Result<f64, AstroError> {
169        // Build the rotation matrix using Orbit Dual.
170        let mut rotmat = SMatrix::<f64, 1, 6>::zeros();
171        let orbit_dual = OrbitGrad::from(self.nominal_state.orbit);
172
173        let xf_partial = orbit_dual.partial_for(param).context(AstroPhysicsSnafu)?;
174        for (cno, val) in [
175            xf_partial.wrt_x(),
176            xf_partial.wrt_y(),
177            xf_partial.wrt_z(),
178            xf_partial.wrt_vx(),
179            xf_partial.wrt_vy(),
180            xf_partial.wrt_vz(),
181        ]
182        .iter()
183        .copied()
184        .enumerate()
185        {
186            rotmat[(0, cno)] = val;
187        }
188
189        Ok((rotmat * self.covar.fixed_view::<6, 6>(0, 0) * rotmat.transpose())[(0, 0)].sqrt())
190    }
191
192    /// Returns the 6x6 covariance (i.e. square of the sigma/uncertainty) of the SMA, ECC, INC, RAAN, AOP, and True Anomaly.
193    pub fn keplerian_covar(&self) -> SMatrix<f64, 6, 6> {
194        // Build the rotation matrix using Orbit Dual.
195        let mut rotmat = SMatrix::<f64, 6, 6>::zeros();
196        let orbit_dual = OrbitGrad::from(self.nominal_state.orbit);
197        for (pno, param) in [
198            OrbitalElement::SemiMajorAxis,
199            OrbitalElement::Eccentricity,
200            OrbitalElement::Inclination,
201            OrbitalElement::RAAN,
202            OrbitalElement::AoP,
203            OrbitalElement::TrueAnomaly,
204        ]
205        .iter()
206        .copied()
207        .enumerate()
208        {
209            let xf_partial = orbit_dual.partial_for(param).unwrap();
210            for (cno, val) in [
211                xf_partial.wrt_x(),
212                xf_partial.wrt_y(),
213                xf_partial.wrt_z(),
214                xf_partial.wrt_vx(),
215                xf_partial.wrt_vy(),
216                xf_partial.wrt_vz(),
217            ]
218            .iter()
219            .copied()
220            .enumerate()
221            {
222                rotmat[(pno, cno)] = val;
223            }
224        }
225
226        rotmat * self.covar.fixed_view::<6, 6>(0, 0) * rotmat.transpose()
227    }
228}
229
230impl<T: State> Estimate<T> for KfEstimate<T>
231where
232    DefaultAllocator: Allocator<<T as State>::Size>
233        + Allocator<<T as State>::Size, <T as State>::Size>
234        + Allocator<<T as State>::Size>
235        + Allocator<<T as State>::VecLength>,
236    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
237    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
238{
239    fn zeros(nominal_state: T) -> Self {
240        Self {
241            nominal_state,
242            state_deviation: OVector::<f64, <T as State>::Size>::zeros(),
243            covar: OMatrix::<f64, <T as State>::Size, <T as State>::Size>::zeros(),
244            covar_bar: OMatrix::<f64, <T as State>::Size, <T as State>::Size>::zeros(),
245            predicted: true,
246            stm: OMatrix::<f64, <T as State>::Size, <T as State>::Size>::identity(),
247        }
248    }
249
250    fn nominal_state(&self) -> T {
251        self.nominal_state
252    }
253
254    fn state_deviation(&self) -> OVector<f64, <T as State>::Size> {
255        self.state_deviation
256    }
257
258    fn covar(&self) -> OMatrix<f64, <T as State>::Size, <T as State>::Size> {
259        self.covar
260    }
261
262    fn predicted_covar(&self) -> OMatrix<f64, <T as State>::Size, <T as State>::Size> {
263        self.covar_bar
264    }
265
266    fn predicted(&self) -> bool {
267        self.predicted
268    }
269    fn stm(&self) -> &OMatrix<f64, <T as State>::Size, <T as State>::Size> {
270        &self.stm
271    }
272    fn set_state_deviation(&mut self, new_state: OVector<f64, <T as State>::Size>) {
273        self.state_deviation = new_state;
274    }
275    fn set_covar(&mut self, new_covar: OMatrix<f64, <T as State>::Size, <T as State>::Size>) {
276        self.covar = new_covar;
277    }
278}
279
280impl<T: State> fmt::Display for KfEstimate<T>
281where
282    DefaultAllocator: Allocator<<T as State>::Size>
283        + Allocator<<T as State>::Size, <T as State>::Size>
284        + Allocator<<T as State>::Size>
285        + Allocator<<T as State>::VecLength>,
286    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
287    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
288{
289    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
290        let dim = <T as State>::Size::dim();
291        let word = if self.predicted {
292            "Prediction"
293        } else {
294            "Estimate"
295        };
296        let mut fmt_cov = Vec::with_capacity(dim);
297        for i in 0..dim {
298            let unit = if i < 3 {
299                "km"
300            } else if i < 6 {
301                "km/s"
302            } else {
303                ""
304            };
305            fmt_cov.push(format!("{:.6} {unit}", &self.covar[(i, i)]));
306        }
307        write!(
308            f,
309            "=== {} @ {} -- within 3 sigma: {} ===\nstate {}\nsigmas [{}]\n",
310            word,
311            &self.epoch(),
312            self.within_3sigma(),
313            &self.state(),
314            fmt_cov.join(", ")
315        )
316    }
317}
318
319impl<T: State> fmt::LowerExp for KfEstimate<T>
320where
321    DefaultAllocator: Allocator<<T as State>::Size>
322        + Allocator<<T as State>::Size, <T as State>::Size>
323        + Allocator<<T as State>::Size>
324        + Allocator<<T as State>::VecLength>,
325    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
326    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
327{
328    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
329        let dim = <T as State>::Size::dim();
330        let word = if self.predicted {
331            "Prediction"
332        } else {
333            "Estimate"
334        };
335        let mut fmt_cov = Vec::with_capacity(dim);
336        for i in 0..dim {
337            let unit = if i < 3 {
338                "km"
339            } else if i < 6 {
340                "km/s"
341            } else {
342                ""
343            };
344            fmt_cov.push(format!("{:e} {unit}", &self.covar[(i, i)]));
345        }
346        write!(
347            f,
348            "=== {} @ {} -- within 3 sigma: {} ===\nstate {}\nsigmas [{}]\n",
349            word,
350            &self.epoch(),
351            self.within_3sigma(),
352            &self.state(),
353            fmt_cov.join(", ")
354        )
355    }
356}
357
358impl<T: State> Mul<f64> for KfEstimate<T>
359where
360    DefaultAllocator: Allocator<<T as State>::Size>
361        + Allocator<<T as State>::Size, <T as State>::Size>
362        + Allocator<<T as State>::Size>
363        + Allocator<<T as State>::VecLength>,
364    <DefaultAllocator as Allocator<<T as State>::Size>>::Buffer<f64>: Copy,
365    <DefaultAllocator as Allocator<<T as State>::Size, <T as State>::Size>>::Buffer<f64>: Copy,
366{
367    type Output = Self;
368
369    fn mul(mut self, rhs: f64) -> Self::Output {
370        self.covar *= rhs.powi(2);
371        self
372    }
373}
374
375#[cfg(test)]
376mod ut_kfest {
377    use crate::{
378        mc::StateDispersion, md::StateParameter, od::estimate::KfEstimate, Spacecraft,
379        GMAT_EARTH_GM,
380    };
381    use anise::analysis::prelude::OrbitalElement;
382    use anise::{constants::frames::EARTH_J2000, prelude::Orbit};
383    use hifitime::Epoch;
384
385    #[test]
386    fn test_estimate_from_disp() {
387        let eme2k = EARTH_J2000.with_mu_km3_s2(GMAT_EARTH_GM);
388        let dt = Epoch::from_gregorian_tai_at_midnight(2020, 1, 1);
389        let initial_state = Spacecraft::builder()
390            .orbit(Orbit::keplerian(
391                22000.0, 0.01, 30.0, 80.0, 40.0, 0.0, dt, eme2k,
392            ))
393            .build();
394
395        let initial_estimate = KfEstimate::disperse_from_diag(
396            initial_state,
397            vec![
398                StateDispersion::builder()
399                    .param(StateParameter::Element(OrbitalElement::SemiMajorAxis))
400                    .std_dev(1.1)
401                    .build(),
402                StateDispersion::zero_mean(
403                    StateParameter::Element(OrbitalElement::Inclination),
404                    0.2,
405                ),
406                StateDispersion::zero_mean(StateParameter::Element(OrbitalElement::RAAN), 0.2),
407                StateDispersion::zero_mean(StateParameter::Element(OrbitalElement::AoP), 0.2),
408            ],
409            Some(0),
410        )
411        .unwrap();
412
413        let initial_state_dev = initial_estimate.nominal_state;
414
415        let (init_rss_pos_km, init_rss_vel_km_s, _) =
416            initial_state.rss(&initial_state_dev).unwrap();
417
418        let delta = (initial_state.orbit - initial_state_dev.orbit).unwrap();
419
420        println!("Truth initial state:\n{initial_state}\n{initial_state:x}");
421        println!("Filter initial state:\n{initial_state_dev}\n{initial_state_dev:x}");
422        println!(
423            "Initial state dev:\t{init_rss_pos_km:.6} km\t{init_rss_vel_km_s:.6} km/s\n{delta}",
424        );
425        println!("covariance: {:.6}", initial_estimate.covar);
426
427        // Check that the error is in the square root of the covariance
428        assert!(delta.radius_km.x < initial_estimate.covar[(0, 0)].sqrt());
429        assert!(delta.radius_km.y < initial_estimate.covar[(1, 1)].sqrt());
430        assert!(delta.radius_km.z < initial_estimate.covar[(2, 2)].sqrt());
431        assert!(delta.velocity_km_s.x < initial_estimate.covar[(3, 3)].sqrt());
432        assert!(delta.velocity_km_s.y < initial_estimate.covar[(4, 4)].sqrt());
433        assert!(delta.velocity_km_s.z < initial_estimate.covar[(5, 5)].sqrt());
434    }
435}
nyx_space/od/estimate/kfestimate.rs

nyx_space/od/estimate/
kfestimate.rs
