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