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