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