Skip to main content

nyx_space/io/
gravity.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 crate::NyxError;
20use crate::linalg::DMatrix;
21use anise::errors::AlmanacError;
22use anise::frames::{Frame, FrameUid};
23use anise::prelude::Almanac;
24use flate2::read::GzDecoder;
25use log::{info, warn};
26use serde::{Deserialize, Serialize};
27use serde_dhall::{SimpleType, StaticType};
28use std::collections::HashMap;
29use std::fmt::Debug;
30use std::fs::File;
31use std::io::Read;
32use std::path::{Path, PathBuf};
33use std::str::FromStr;
34
35/// Unnormalized J2 value of Earth
36pub const EARTH_J2: f64 = -0.484169548456e-03;
37
38#[cfg(feature = "python")]
39use pyo3::prelude::*;
40
41/// Configuration holder for gravity field.
42///
43/// Data is first loaded as a SHADR, if that fails, Nyx will try to load it as a COF file.
44#[derive(Clone, Serialize, Deserialize, Debug)]
45#[cfg_attr(feature = "python", pyclass(from_py_object, get_all, set_all))]
46pub struct GravityFieldConfig {
47    /// Desired degree
48    pub degree: usize,
49    /// Desired order
50    pub order: usize,
51    /// Path to the file, relative to the current working directory
52    pub filepath: PathBuf,
53    /// Set to true if the data is gunzipped
54    pub gunzipped: bool,
55    /// The frame in which to compute this gravity field
56    pub frame: FrameUid,
57}
58
59#[cfg(feature = "python")]
60#[cfg_attr(feature = "python", pymethods)]
61impl GravityFieldConfig {
62    #[pyo3(signature=(degree, order, filepath, frame, gunzipped=true))]
63    #[new]
64    fn py_new(
65        degree: usize,
66        order: usize,
67        filepath: PathBuf,
68        frame: FrameUid,
69        gunzipped: bool,
70    ) -> Self {
71        Self {
72            filepath,
73            gunzipped,
74            degree,
75            order,
76            frame,
77        }
78    }
79
80    fn __str__(&self) -> String {
81        format!("{self:?}")
82    }
83
84    fn __repr__(&self) -> String {
85        format!("{self:?} @ {self:p}")
86    }
87}
88
89/// `GravityFieldData` loads the requested gravity potential files and stores them in memory (in a HashMap).
90///
91/// WARNING: This memory backend may require a lot of RAM (e.g. EMG2008 2190x2190 requires nearly 400 MB of RAM).
92#[derive(Clone)]
93pub struct GravityFieldData {
94    degree: usize,
95    order: usize,
96    c_nm: DMatrix<f64>,
97    s_nm: DMatrix<f64>,
98    pub frame: Frame,
99}
100
101impl GravityFieldData {
102    pub fn from_config(cfg: GravityFieldConfig, almanac: &Almanac) -> Result<Self, NyxError> {
103        let frame = almanac
104            .frame_info(cfg.frame)
105            .map_err(|e| NyxError::FromAlmanacError {
106                source: Box::new(AlmanacError::GenericError { err: e.to_string() }),
107                action: "fetching gravity field frame",
108            })?;
109
110        if !cfg.gunzipped && cfg.filepath.ends_with(".cof")
111            || cfg.gunzipped && cfg.filepath.ends_with(".cof.gz")
112        {
113            Self::from_cof(cfg.filepath, cfg.degree, cfg.order, cfg.gunzipped, frame)
114        } else {
115            Self::from_shadr(cfg.filepath, cfg.degree, cfg.order, cfg.gunzipped, frame)
116        }
117    }
118
119    /// Initialize `GravityFieldData` with a custom unnormalized J2 value
120    pub fn from_j2(j2: f64, frame: Frame) -> GravityFieldData {
121        let mut c_nm = DMatrix::from_element(3, 3, 0.0);
122        c_nm[(2, 0)] = j2;
123
124        GravityFieldData {
125            degree: 2,
126            order: 0,
127            c_nm,
128            s_nm: DMatrix::from_element(3, 3, 0.0),
129            frame,
130        }
131    }
132
133    /// Initialize `GravityFieldData` from the file path (must be a gunzipped file)
134    ///
135    /// Gravity models provided by `nyx`:
136    /// + EMG2008 to 2190 for Earth (tide free)
137    /// + Moon to 1500 (from SHADR file)
138    /// + Mars to 120 (from SHADR file)
139    /// + Venus to 150 (from SHADR file)
140    pub fn from_shadr<P: AsRef<Path> + Debug>(
141        filepath: P,
142        degree: usize,
143        order: usize,
144        gunzipped: bool,
145        frame: Frame,
146    ) -> Result<GravityFieldData, NyxError> {
147        Self::load(
148            filepath, gunzipped, true, //SHADR has a header which we ignore
149            degree, order, frame,
150        )
151    }
152
153    pub fn from_cof<P: AsRef<Path> + Debug>(
154        filepath: P,
155        degree: usize,
156        order: usize,
157        gunzipped: bool,
158        frame: Frame,
159    ) -> Result<GravityFieldData, NyxError> {
160        let mut f = File::open(&filepath).map_err(|_| NyxError::FileUnreadable {
161            msg: format!("File not found: {filepath:?}"),
162        })?;
163        let mut buffer = vec![0; 0];
164        if gunzipped {
165            let mut d = GzDecoder::new(f);
166            d.read_to_end(&mut buffer)
167                .map_err(|_| NyxError::FileUnreadable {
168                    msg: "could not read file as gunzip".to_string(),
169                })?;
170        } else {
171            f.read_to_end(&mut buffer)
172                .map_err(|_| NyxError::FileUnreadable {
173                    msg: "could not read file to end".to_string(),
174                })?;
175        }
176
177        let data_as_str = String::from_utf8(buffer).map_err(|_| NyxError::FileUnreadable {
178            msg: "could not decode file contents as utf8".to_string(),
179        })?;
180
181        // Since the COF files are so specific, we just code everything up in here.
182
183        let mut c_nm_mat = DMatrix::from_element(degree + 1, degree + 1, 0.0);
184        let mut s_nm_mat = DMatrix::from_element(degree + 1, degree + 1, 0.0);
185        let mut max_order: usize = 0;
186        let mut max_degree: usize = 0;
187        for (lno, line) in data_as_str.split('\n').enumerate() {
188            if line.is_empty() || !line.starts_with('R') {
189                continue; // This is either a comment, a header or "END"
190            }
191            // These variables need to be declared as mutable because rustc does not know
192            // we nwon't match each ino more than once.
193            let mut cur_degree: usize = 0;
194            let mut cur_order: usize = 0;
195            let mut c_nm: f64 = 0.0;
196            let mut s_nm: f64 = 0.0;
197            for (ino, item) in line.split_whitespace().enumerate() {
198                match ino {
199                    0 => continue, // We need this so we don't break at every first item
200                    1 => match usize::from_str(item) {
201                        Ok(val) => cur_degree = val,
202                        Err(_) => {
203                            return Err(NyxError::FileUnreadable {
204                                msg: format!(
205                                    "Harmonics file:
206                                could not parse degree `{item}` on line {lno}"
207                                ),
208                            });
209                        }
210                    },
211                    2 => match usize::from_str(item) {
212                        Ok(val) => cur_order = val,
213                        Err(_) => {
214                            return Err(NyxError::FileUnreadable {
215                                msg: format!(
216                                    "Harmonics file:
217                                could not parse order `{item}` on line {lno}"
218                                ),
219                            });
220                        }
221                    },
222                    3 => {
223                        // If we are at degree zero, then there is only one item, so we can parse that and
224                        // set the S_nm to zero.
225                        if degree == 0 {
226                            s_nm = 0.0;
227                            match f64::from_str(item) {
228                                Ok(val) => c_nm = val,
229                                Err(_) => {
230                                    return Err(NyxError::FileUnreadable {
231                                        msg: format!(
232                                            "Harmonics file:
233                                        could not parse C_nm `{item}` on line {lno}"
234                                        ),
235                                    });
236                                }
237                            }
238                        } else {
239                            // There is a space as a delimiting character between the C_nm and S_nm only if the S_nm
240                            // is a positive number, otherwise, they are continuous (what a great format).
241                            if (item.matches('-').count() == 3 && !item.starts_with('-'))
242                                || item.matches('-').count() == 4
243                            {
244                                // Now we have two items concatenated into one... great
245                                let parts: Vec<&str> = item.split('-').collect();
246                                if parts.len() == 5 {
247                                    // That mean we have five minus signs, so both the C and S are negative.
248                                    let c_nm_str = "-".to_owned() + parts[1] + "-" + parts[2];
249                                    match f64::from_str(&c_nm_str) {
250                                        Ok(val) => c_nm = val,
251                                        Err(_) => {
252                                            return Err(NyxError::FileUnreadable {
253                                                msg: format!(
254                                                    "Harmonics file:
255                                                could not parse C_nm `{item}` on line {lno}"
256                                                ),
257                                            });
258                                        }
259                                    }
260                                    // That mean we have five minus signs, so both the C and S are negative.
261                                    let s_nm_str = "-".to_owned() + parts[3] + "-" + parts[4];
262                                    match f64::from_str(&s_nm_str) {
263                                        Ok(val) => s_nm = val,
264                                        Err(_) => {
265                                            return Err(NyxError::FileUnreadable {
266                                                msg: format!(
267                                                    "Harmonics file:
268                                                could not parse S_nm `{item}` on line {lno}"
269                                                ),
270                                            });
271                                        }
272                                    }
273                                } else {
274                                    // That mean we have fouve minus signs, and since both values are concatenated, C_nm is positive and S_nm is negative
275                                    let c_nm_str = parts[0].to_owned() + "-" + parts[1];
276                                    match f64::from_str(&c_nm_str) {
277                                        Ok(val) => c_nm = val,
278                                        Err(_) => {
279                                            return Err(NyxError::FileUnreadable {
280                                                msg: format!(
281                                                    "Harmonics file:
282                                                could not parse C_nm `{item}` on line {lno}"
283                                                ),
284                                            });
285                                        }
286                                    }
287                                    // That mean we have five minus signs, so both the C and S are negative.
288                                    let s_nm_str = "-".to_owned() + parts[2] + "-" + parts[3];
289                                    match f64::from_str(&s_nm_str) {
290                                        Ok(val) => s_nm = val,
291                                        Err(_) => {
292                                            return Err(NyxError::FileUnreadable {
293                                                msg: format!(
294                                                    "Harmonics file:
295                                                could not parse S_nm `{item}` on line {lno}"
296                                                ),
297                                            });
298                                        }
299                                    }
300                                }
301                            } else {
302                                // We only have the first item, and that's the C_nm
303                                match f64::from_str(item) {
304                                    Ok(val) => c_nm = val,
305                                    Err(_) => {
306                                        return Err(NyxError::FileUnreadable {
307                                            msg: format!(
308                                                "Harmonics file:
309                                            could not parse C_nm `{item}` on line {lno}"
310                                            ),
311                                        });
312                                    }
313                                }
314                            }
315                        }
316                    }
317                    4 => match f64::from_str(item) {
318                        // If this exists, then the S_nm is positive.
319                        Ok(val) => s_nm = val,
320                        Err(_) => {
321                            return Err(NyxError::FileUnreadable {
322                                msg: format!(
323                                    "Harmonics file:
324                                could not parse S_nm `{item}` on line {lno}"
325                                ),
326                            });
327                        }
328                    },
329                    _ => break, // We aren't storing the covariance of these harmonics
330                }
331            }
332
333            if cur_degree > degree {
334                // The file is organized by degree, so once we've passed the maximum degree we want,
335                // we can safely stop reading the file.
336                break;
337            }
338
339            // Only insert this data into the hashmap if it's within the required order as well
340            if cur_order <= order {
341                c_nm_mat[(cur_degree, cur_order)] = c_nm;
342                s_nm_mat[(cur_degree, cur_order)] = s_nm;
343            }
344            // This serves as a warning.
345            max_order = if cur_order > max_order {
346                cur_order
347            } else {
348                max_order
349            };
350            max_degree = if cur_degree > max_degree {
351                cur_degree
352            } else {
353                max_degree
354            };
355        }
356        if max_degree < degree || max_order < order {
357            warn!(
358                "{filepath:?} only contained (degree, order) of ({max_degree}, {max_order}) instead of requested ({degree}, {order})"
359            );
360        } else {
361            info!("{filepath:?} loaded with (degree, order) = ({degree}, {order})");
362        }
363        Ok(GravityFieldData {
364            degree: max_degree,
365            order: max_order,
366            c_nm: c_nm_mat,
367            s_nm: s_nm_mat,
368            frame,
369        })
370    }
371
372    /// `load` handles the actual loading in memory.
373    fn load<P: AsRef<Path> + Debug>(
374        filepath: P,
375        gunzipped: bool,
376        skip_first_line: bool,
377        degree: usize,
378        order: usize,
379        frame: Frame,
380    ) -> Result<GravityFieldData, NyxError> {
381        let mut f = File::open(&filepath).map_err(|_| NyxError::FileUnreadable {
382            msg: format!("File not found: {filepath:?}"),
383        })?;
384        let mut buffer = vec![0; 0];
385        if gunzipped {
386            let mut d = GzDecoder::new(f);
387            d.read_to_end(&mut buffer)
388                .map_err(|_| NyxError::FileUnreadable {
389                    msg: "could not read file as gunzip".to_string(),
390                })?;
391        } else {
392            f.read_to_end(&mut buffer)
393                .map_err(|_| NyxError::FileUnreadable {
394                    msg: "could not read file to end".to_string(),
395                })?;
396        }
397
398        let data_as_str = String::from_utf8(buffer).map_err(|_| NyxError::FileUnreadable {
399            msg: "could not decode file contents as utf8".to_string(),
400        })?;
401
402        let mut c_nm_mat = DMatrix::from_element(degree + 1, degree + 1, 0.0);
403        let mut s_nm_mat = DMatrix::from_element(degree + 1, degree + 1, 0.0);
404
405        let mut max_degree: usize = 0;
406        let mut max_order: usize = 0;
407        for (lno, line) in data_as_str.split('\n').enumerate() {
408            if lno == 0 && skip_first_line {
409                continue;
410            }
411            // These variables need to be declared as mutable because rustc does not know
412            // we won't match each ino more than once.
413            let mut cur_order: usize = 0;
414            let mut cur_degree: usize = 0;
415            let mut c_nm: f64 = 0.0;
416            let mut s_nm: f64 = 0.0;
417            for (ino, item) in line.replace(',', " ").split_whitespace().enumerate() {
418                match ino {
419                    0 => match usize::from_str(item) {
420                        Ok(val) => cur_degree = val,
421                        Err(_) => {
422                            return Err(NyxError::FileUnreadable {
423                                msg: format!(
424                                    "Harmonics file:
425                                could not parse degree on line {lno} (`{item}`)",
426                                ),
427                            });
428                        }
429                    },
430                    1 => match usize::from_str(item) {
431                        Ok(val) => cur_order = val,
432                        Err(_) => {
433                            return Err(NyxError::FileUnreadable {
434                                msg: format!(
435                                    "Harmonics file:
436                                could not parse order on line {lno} (`{item}`)"
437                                ),
438                            });
439                        }
440                    },
441                    2 => match f64::from_str(&item.replace('D', "E")) {
442                        Ok(val) => c_nm = val,
443                        Err(_) => {
444                            return Err(NyxError::FileUnreadable {
445                                msg: format!(
446                                    "Harmonics file:
447                                could not parse C_nm `{item}` on line {lno}"
448                                ),
449                            });
450                        }
451                    },
452                    3 => match f64::from_str(&item.replace('D', "E")) {
453                        Ok(val) => s_nm = val,
454                        Err(_) => {
455                            return Err(NyxError::FileUnreadable {
456                                msg: format!(
457                                    "Harmonics file:
458                                could not parse S_nm `{item}` on line {lno}"
459                                ),
460                            });
461                        }
462                    },
463                    _ => break, // We aren't storing the covariance of these harmonics
464                }
465            }
466
467            if cur_degree > degree {
468                // The file is organized by degree, so once we've passed the maximum degree we want,
469                // we can safely stop reading the file.
470                break;
471            }
472
473            // Only insert this data into the hashmap if it's within the required order as well
474            if cur_order <= order {
475                c_nm_mat[(cur_degree, cur_order)] = c_nm;
476                s_nm_mat[(cur_degree, cur_order)] = s_nm;
477            }
478            // This serves as a warning.
479            max_order = if cur_order > max_order {
480                cur_order
481            } else {
482                max_order
483            };
484            max_degree = if cur_degree > max_degree {
485                cur_degree
486            } else {
487                max_degree
488            };
489        }
490        if max_degree < degree || max_order < order {
491            warn!(
492                "{filepath:?} only contained (degree, order) of ({max_degree}, {max_order}) instead of requested ({degree}, {order})",
493            );
494        } else {
495            info!("{filepath:?} loaded with (degree, order) = ({degree}, {order})");
496        }
497        Ok(GravityFieldData {
498            order: max_order,
499            degree: max_degree,
500            c_nm: c_nm_mat,
501            s_nm: s_nm_mat,
502            frame,
503        })
504    }
505
506    /// Returns the maximum order of this gravity potential storage (Jnm=Jn2,Jn3...)
507    pub fn max_order_m(&self) -> usize {
508        self.order
509    }
510
511    /// Returns the maximum degree of this gravity potential storage (Jn=J2,J3...)
512    pub fn max_degree_n(&self) -> usize {
513        self.degree
514    }
515
516    /// Returns the C_nm and S_nm for the provided order and degree.
517    pub fn cs_nm(&self, degree: usize, order: usize) -> (f64, f64) {
518        (self.c_nm[(degree, order)], self.s_nm[(degree, order)])
519    }
520}
521
522impl StaticType for GravityFieldConfig {
523    fn static_type() -> SimpleType {
524        let mut fields = HashMap::new();
525
526        fields.insert("filepath".to_string(), String::static_type());
527        fields.insert("gunzipped".to_string(), bool::static_type());
528        fields.insert("degree".to_string(), usize::static_type());
529        fields.insert("order".to_string(), usize::static_type());
530
531        SimpleType::Record(fields)
532    }
533}
534
535#[cfg(test)]
536#[test]
537fn test_load_harmonic_files() {
538    use anise::constants::frames::IAU_EARTH_FRAME;
539
540    let data_folder: PathBuf = [env!("CARGO_MANIFEST_DIR"), "../data/01_planetary"]
541        .iter()
542        .collect();
543
544    GravityFieldData::from_cof(
545        data_folder.join("JGM3.cof.gz"),
546        50,
547        50,
548        true,
549        IAU_EARTH_FRAME,
550    )
551    .expect("could not load JGM3");
552
553    GravityFieldData::from_shadr(
554        data_folder.join("EGM2008_to2190_TideFree.gz"),
555        120,
556        120,
557        true,
558        IAU_EARTH_FRAME,
559    )
560    .expect("could not load EGM2008");
561
562    GravityFieldData::from_shadr(
563        data_folder.join("Luna_jggrx_1500e_sha.tab.gz"),
564        1500,
565        1500,
566        true,
567        IAU_EARTH_FRAME,
568    )
569    .expect("could not load jggrx");
570}