Add epoch propagation function

ajtribick · Jun 21, 2021 · 1023824 · 1023824
1 parent af30231
commit 1023824
Show file tree

Hide file tree

Showing 5 changed files with 142 additions and 20 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -17,5 +17,6 @@ byteorder = "1.4"
 flate2 = "1.0"
 globset = "0.4.8"
 memchr = "2.4"
+nalgebra = "0.27.1"
 pyo3 = {version="0.13.2", features=["extension-module"]}
 quick-xml = "0.22.0"
diff --git a/celestia_gaia/gaia_data.py b/celestia_gaia/gaia_data.py
@@ -23,7 +23,7 @@
 
 from astroquery.gaia import Gaia
 
-from .directories import GAIA_EDR3_DIR, VIZIER_DIR
+from .directories import GAIA_EDR3_DIR
 from .ranges import MultiRange
 from .utils import confirm_action
 from .celestia_gaia import check_hip_ids

diff --git a/src/astro.rs b/src/astro.rs
@@ -0,0 +1,91 @@
+use std::f64;
+
+use nalgebra::Vector3;
+
+const MAS_TO_RADIANS: f64 = f64::consts::PI / (3600000.0 * 180.0);
+const AU_IN_KM_YEAR_PER_S: f64 = 149597870.7 / (365.25 * 86400.0);
+
+// basic linear motion, not accounting for light travel time
+pub fn apply_pm(
+    ra: f64,
+    dec: f64,
+    pm_ra: f64,
+    pm_dec: f64,
+    rv: f64,
+    parallax: f64,
+    epoch1: f64,
+    epoch2: f64,
+) -> (f64, f64) {
+    let (ra_sin, ra_cos) = ra.to_radians().sin_cos();
+    let (dec_sin, dec_cos) = dec.to_radians().sin_cos();
+    let start = Vector3::new(ra_cos * dec_cos, ra_sin * dec_cos, dec_sin);
+
+    // get coordinate vectors
+    let z = start.normalize();
+    let x = Vector3::new(-ra_sin, ra_cos, 0.0);
+    let y = z.cross(&x);
+
+    let velocity = pm_ra * MAS_TO_RADIANS * x
+        + pm_dec * MAS_TO_RADIANS * y
+        + rv * parallax * MAS_TO_RADIANS / AU_IN_KM_YEAR_PER_S * z;
+
+    let end = start + velocity * (epoch2 - epoch1);
+
+    (
+        end.y.atan2(end.x).to_degrees(),
+        (end.z / end.norm()).asin().to_degrees(),
+    )
+}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    // test cases from ADQL EPOCH_PROP_POS in Gaia archive
+
+    #[test]
+    fn apply_pm_zero_rv() {
+        let (ra, dec, parallax, pm_ra, pm_dec, rv, epoch1, epoch2) =
+            (132.0, -38.0, 850.0, 4200.0, 3200.0, 0.0, 1990.0, 2020.0);
+        let ra_expected = 2.30460952981089_f64.to_degrees();
+        let dec_expected = -0.662759549026617_f64.to_degrees();
+        let (ra_actual, dec_actual) =
+            apply_pm(ra, dec, pm_ra, pm_dec, rv, parallax, epoch1, epoch2);
+        let ra_diff = (ra_actual - ra_expected).abs();
+        let dec_diff = (dec_actual - dec_expected).abs();
+        let max_diff = f64::max(ra_diff, dec_diff);
+        assert!(
+            max_diff < 1e-9,
+            "Coordinate mismatch: expected ({}, {}), actual ({}, {}), diff ({}, {})",
+            ra_expected,
+            dec_expected,
+            ra_actual,
+            dec_actual,
+            ra_diff,
+            dec_diff
+        );
+    }
+
+    #[test]
+    fn apply_pm_with_rv() {
+        let (ra, dec, parallax, pm_ra, pm_dec, rv, epoch1, epoch2) =
+            (132.0, -38.0, 850.0, 4200.0, 3200.0, 80.0, 1990.0, 2020.0);
+        let ra_expected = 2.30460791702764_f64.to_degrees();
+        let dec_expected = -0.662760518633619_f64.to_degrees();
+        let (ra_actual, dec_actual) =
+            apply_pm(ra, dec, pm_ra, pm_dec, rv, parallax, epoch1, epoch2);
+        let ra_diff = (ra_actual - ra_expected).abs();
+        let dec_diff = (dec_actual - dec_expected).abs();
+        let max_diff = f64::max(ra_diff, dec_diff);
+        assert!(
+            max_diff < 1e-9,
+            "Coordinate mismatch: expected ({}, {}), actual ({}, {}), diff ({}, {})",
+            ra_expected,
+            dec_expected,
+            ra_actual,
+            dec_actual,
+            ra_diff,
+            dec_diff
+        );
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
@@ -6,6 +6,7 @@ use std::{
 use globset::Glob;
 use pyo3::{prelude::*, wrap_pyfunction};
 
+mod astro;
 mod error;
 mod votable;
 
@@ -14,7 +15,9 @@ use votable::VotableReader;
 
 #[pyfunction]
 fn check_hip_ids(gaia_dir: String) -> PyResult<()> {
-    let pattern = Glob::new("**/gaiaedr3-hip2-*.vot.gz").map_err(Error::new)?.compile_matcher();
+    let pattern = Glob::new("**/gaiaedr3-hip2-*.vot.gz")
+        .map_err(Error::new)?
+        .compile_matcher();
     let mut hip_to_gaia: HashMap<i32, Vec<i64>> = HashMap::new();
     let mut gaia_to_hip: HashMap<i64, Vec<i32>> = HashMap::new();
     for entry_result in read_dir(gaia_dir)? {
@@ -39,8 +42,14 @@ fn check_hip_ids(gaia_dir: String) -> PyResult<()> {
         while let Some(accessor) = reader.read()? {
             let hip = accessor.read_i32(hip_ordinal)?.unwrap();
             let source_id = accessor.read_i64(gaia_ordinal)?.unwrap();
-            hip_to_gaia.entry(hip).and_modify(|v| v.push(source_id)).or_insert(vec![source_id]);
-            gaia_to_hip.entry(source_id).and_modify(|v| v.push(hip)).or_insert(vec![hip]);
+            hip_to_gaia
+                .entry(hip)
+                .and_modify(|v| v.push(source_id))
+                .or_insert(vec![source_id]);
+            gaia_to_hip
+                .entry(source_id)
+                .and_modify(|v| v.push(hip))
+                .or_insert(vec![hip]);
         }
     }
 

diff --git a/src/votable.rs b/src/votable.rs
@@ -1,5 +1,4 @@
 use std::{
-    borrow::Cow,
     cmp,
     collections::HashMap,
     io::{self, BufRead, BufReader, ErrorKind, Read},
@@ -56,9 +55,7 @@ fn parse_field(attributes: Attributes) -> Result<(Vec<u8>, DataType), Error> {
         match attribute.key {
             b"name" => name = Some(attribute.value.into_owned()),
             b"datatype" => datatype = Some(DataType::parse_bytes(&attribute.value)?),
-            b"arraysize" => {
-                return Err(Error::parse("Array types not supported"))
-            }
+            b"arraysize" => return Err(Error::parse("Array types not supported")),
             _ => (),
         }
     }
@@ -151,7 +148,7 @@ impl<R: Read> VotableReader<R> {
         Ok(Some(RecordAccessor {
             mask: (&self.buffer[..self.mask_width]).view_bits(),
             field_types: &self.field_types,
-            field_offsets: Cow::from(&self.field_offsets),
+            field_offsets: &self.field_offsets,
             data: &self.buffer[self.mask_width..],
         }))
     }
@@ -160,7 +157,7 @@ impl<R: Read> VotableReader<R> {
 pub struct RecordAccessor<'a> {
     mask: &'a BitSlice<Msb0, u8>,
     field_types: &'a [DataType],
-    field_offsets: Cow<'a, [usize]>,
+    field_offsets: &'a [usize],
     data: &'a [u8],
 }
 
@@ -175,8 +172,9 @@ impl<'a> RecordAccessor<'a> {
         }
 
         let offset = self.field_offsets[ordinal];
-        Ok(Some((&self.data[offset..offset + std::mem::size_of::<i16>()])
-            .read_i16::<BigEndian>()?))
+        Ok(Some(
+            (&self.data[offset..offset + std::mem::size_of::<i16>()]).read_i16::<BigEndian>()?,
+        ))
     }
 
     pub fn read_i32(&self, ordinal: usize) -> Result<Option<i32>, Error> {
@@ -189,8 +187,9 @@ impl<'a> RecordAccessor<'a> {
         }
 
         let offset = self.field_offsets[ordinal];
-        Ok(Some((&self.data[offset..offset + std::mem::size_of::<i32>()])
-            .read_i32::<BigEndian>()?))
+        Ok(Some(
+            (&self.data[offset..offset + std::mem::size_of::<i32>()]).read_i32::<BigEndian>()?,
+        ))
     }
 
     pub fn read_i64(&self, ordinal: usize) -> Result<Option<i64>, Error> {
@@ -203,8 +202,9 @@ impl<'a> RecordAccessor<'a> {
         }
 
         let offset = self.field_offsets[ordinal];
-        Ok(Some((&self.data[offset..offset + std::mem::size_of::<i64>()])
-            .read_i64::<BigEndian>()?))
+        Ok(Some(
+            (&self.data[offset..offset + std::mem::size_of::<i64>()]).read_i64::<BigEndian>()?,
+        ))
     }
 
     pub fn read_f32(&self, ordinal: usize) -> Result<Option<f32>, Error> {
@@ -217,8 +217,9 @@ impl<'a> RecordAccessor<'a> {
         }
 
         let offset = self.field_offsets[ordinal];
-        Ok(Some((&self.data[offset..offset + std::mem::size_of::<f32>()])
-            .read_f32::<BigEndian>()?))
+        Ok(Some(
+            (&self.data[offset..offset + std::mem::size_of::<f32>()]).read_f32::<BigEndian>()?,
+        ))
     }
 
     pub fn read_f64(&self, ordinal: usize) -> Result<Option<f64>, Error> {
@@ -231,8 +232,9 @@ impl<'a> RecordAccessor<'a> {
         }
 
         let offset = self.field_offsets[ordinal];
-        Ok(Some((&self.data[offset..offset + std::mem::size_of::<f64>()])
-            .read_f64::<BigEndian>()?))
+        Ok(Some(
+            (&self.data[offset..offset + std::mem::size_of::<f64>()]).read_f64::<BigEndian>()?,
+        ))
     }
 }
 
@@ -313,3 +315,22 @@ impl<R: BufRead> BufRead for Binary2Reader<R> {
         self.position += amt;
     }
 }
+
+#[cfg(test)]
+mod test {
+    use super::*;
+
+    #[test]
+    fn binary2_read() {
+        let source: &[u8] = b"AgMFBwsNERMXHR8l\nKSsvNTs9Q0c=\n</STREAM>";
+        let expected = [
+            2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71,
+        ];
+        let buf_reader = BufReader::with_capacity(5, source);
+        let mut reader = Binary2Reader::new(buf_reader);
+        let mut actual = Vec::new();
+        let length = reader.read_to_end(&mut actual).unwrap();
+        assert_eq!(length, 20);
+        assert_eq!(&expected, actual.as_slice());
+    }
+}