Merge pull request #319 from sblunt/hipiad-bugfixes

add hip gaia save/load unit test

orbitize/gaia.py

@@ -28,10 +28,22 @@ class currently only fits for the position of the star in the Gaia epoch,
         hiplogprob (orbitize.hipparcos.HipLogProb): object containing
             all info relevant to Hipparcos IAD fitting
         dr (str): either 'dr2' or 'edr3'
+        query (bool): if True, queries the Gaia database for astrometry of the
+            target (requires an internet connection). If False, uses user-input 
+            astrometric values (runs without internet).
+        gaia_data (dict): see `query` keyword above. If `query` set to False, 
+            then user must supply a dictionary of Gaia astometry in the following
+            form:
+                gaia_data = {
+                    'ra': 139.4 # RA in degrees
+                    'dec': 139.4 # Dec in degrees
+                    'ra_error': 0.004 # RA error in mas
+                    'dec_error': 0.004 # Dec error in mas
+                }
 
     Written: Sarah Blunt, 2021
     """
-    def __init__(self, gaia_num, hiplogprob, dr='dr2'):
+    def __init__(self, gaia_num, hiplogprob, dr='dr2', query=True, gaia_data=None):
 
         self.gaia_num = gaia_num
         self.hiplogprob = hiplogprob
@@ -46,15 +58,16 @@ def __init__(self, gaia_num, hiplogprob, dr='dr2'):
         self.hipparcos_epoch = 1991.25
 
 
-        query = """SELECT
-        TOP 1
-        ra, dec, ra_error, dec_error
-        FROM gaia{}.gaia_source
-        WHERE source_id = {}
-        """.format(self.dr, self.gaia_num)
+        if query:
+            query = """SELECT
+            TOP 1
+            ra, dec, ra_error, dec_error
+            FROM gaia{}.gaia_source
+            WHERE source_id = {}
+            """.format(self.dr, self.gaia_num)
 
-        job = Gaia.launch_job_async(query)
-        gaia_data = job.get_results()
+            job = Gaia.launch_job_async(query)
+            gaia_data = job.get_results()
 
         self.ra = gaia_data['ra']
         self.ra_err = gaia_data['ra_error']

orbitize/hipparcos.py

@@ -5,17 +5,17 @@
 from scipy.stats import norm
 import matplotlib.pyplot as plt
 
-from astroquery.vizier import Vizier
 from astropy.time import Time
 from astropy.coordinates import get_body_barycentric_posvel
+from astroquery.vizier import Vizier
 
 class HipparcosLogProb(object):
     """
     Class to compute the log probability of an orbit with respect to the 
-    Hipparcos Intermediate Astrometric Data (IAD). Queries Vizier for 
-    all metadata relevant to the IAD, and reads in the IAD themselves from
-    a specified location. Follows Nielsen+ 2020 (studying the orbit of beta 
-    Pic b).
+    Hipparcos Intermediate Astrometric Data (IAD). If using a DVD file, 
+    queries Vizier for all metadata relevant to the IAD, and reads in the IAD 
+    themselves from a specified location. Follows Nielsen+ 2020 (studying the 
+    orbit of beta Pic b).
 
     Fitting the Hipparcos IAD requires fitting for the following five parameters.
     They are added to the vector of fitting parameters in system.py, but 

tests/test_gaia.py

@@ -32,7 +32,7 @@ def test_system_setup():
     Test that a System object with Hipparcos and Gaia is initialized correctly
     """
     hip_num = '027321' # beta Pic
-    edr3_num = '4792774797545800832'
+    edr3_num = 4792774797545800832
     num_secondary_bodies = 1
     path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
 
@@ -61,13 +61,12 @@ def test_system_setup():
     assert betaPic_system.fit_secondary_mass
     assert betaPic_system.track_planet_perturbs
 
-
 def test_valueerror():
     """
     Check that if I don't say dr2 or edr3, I get a value error
     """
     hip_num = '027321' # beta Pic
-    edr3_num = '4792774797545800832'
+    edr3_num = 4792774797545800832
     num_secondary_bodies = 1
     path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
 
@@ -112,7 +111,7 @@ def test_orbit_calculation():
     d0 = 0
 
     hip_num = '027321' # beta Pic
-    edr3_num = '4792774797545800832'
+    edr3_num = 4792774797545800832
     num_secondary_bodies = 1
     path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
 
@@ -198,8 +197,27 @@ def test_orbit_calculation():
 
     assert np.isclose(np.exp(lnlike), 1)
 
+def test_nointernet():
+    """
+    Test that the internet-less object setup works
+    """
+    hip_num = '027321' # beta Pic
+    dr2_number = 4792774797545105664
+
+    num_secondary_bodies = 1
+    path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
+
+    myHip = hipparcos.HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies)
+
+    dr3Gaia = gaia.GaiaLogProb(
+        dr2_number, myHip, dr='dr2', query=False, gaia_data = {'ra':0, 'dec':0, 'ra_error':0, 'dec_error':0}
+    )
+
+
+
 if __name__ == '__main__':
-    test_dr2_edr3()
+    test_nointernet()
+    # test_dr2_edr3()
     # test_system_setup()
     # test_valueerror()
     # test_orbit_calculation()

tests/test_hipparcos.py

@@ -3,7 +3,8 @@
 
 import matplotlib.pyplot as plt
 
-from orbitize import DATADIR, read_input, system
+from orbitize import DATADIR, read_input, system, sampler, results
+from orbitize.gaia import GaiaLogProb
 from orbitize.hipparcos import HipparcosLogProb, nielsen_iad_refitting_test
 
 def test_hipparcos_api():
@@ -71,7 +72,6 @@ def test_hipparcos_api():
     rejected_scansHip = HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies)
     assert len(rejected_scansHip.cos_phi) == len(raw_iad_data[0]) - 1
 
-
 def test_dvd_vs_2021catalog():
     """
     Test code's ability to parse both a DVD data file and a 2021
@@ -128,9 +128,41 @@ def test_iad_refitting():
     assert np.isclose(0, np.median(post[:, -1]), atol=0.1)
     assert np.isclose(myHipLogProb.plx0, np.median(post[:, 0]), atol=0.1)
 
+def test_save_load():
+    """
+    Set up a Hip IAD + Gaia fit, save the results, and load them.
+    """
+
+    hip_num = '027321' # beta Pic
+
+    num_secondary_bodies = 1
+    path_to_iad_file = '{}HIP{}.d'.format(DATADIR, hip_num)
+
+    myHip = HipparcosLogProb(path_to_iad_file, hip_num, num_secondary_bodies)
+    myGaia = GaiaLogProb(4792774797545800832, myHip, dr='edr3')
+
+    input_file = os.path.join(DATADIR, 'HD4747.csv')
+    data_table_with_rvs = read_input.read_file(input_file)
+    mySys = system.System(
+        1, data_table_with_rvs, 1.22, 56.95, mass_err=0.08, plx_err=0.26, 
+        hipparcos_IAD=myHip, fit_secondary_mass=True, gaia=myGaia
+    )
+    n_walkers = 50
+    mySamp = sampler.MCMC(mySys, num_walkers=n_walkers)
+    mySamp.run_sampler(n_walkers, burn_steps=0)
+    filename = 'tmp.hdf5'
+    mySamp.results.save_results(filename)
+
+    myResults = results.Results()
+    myResults.load_results(filename)
+
+    os.system('rm tmp.hdf5')
+
+
 
 if __name__ == '__main__':
-    test_hipparcos_api()
-    test_iad_refitting()
-    test_dvd_vs_2021catalog()
+    test_save_load()
+    # test_hipparcos_api()
+    # test_iad_refitting()
+    # test_dvd_vs_2021catalog()