Bugfix for relative rvs with massive secondary

orbitize/__init__.py

@@ -1,6 +1,6 @@
 import os
 
-__version__ = "3.4.0"
+__version__ = "4.0.0"
 
 # set Python env variable to keep track of example data dir
 orbitize_dir = os.path.dirname(__file__)

orbitize/kepler.py

@@ -123,9 +123,11 @@ def calc_orbit(
         plx (np.array): parallax [mas]
         mtot (np.array): total mass of the two-body orbit (M_* + M_planet) [Solar masses]
         mass_for_Kamp (np.array, optional): mass of the body that causes the RV signal.
-            For example, if you want to return the stellar RV, this is the planet mass.
-            If you want to return the planetary RV, this is the stellar mass. [Solar masses].
-            For planet mass ~ 0, mass_for_Kamp ~ M_tot, and function returns planetary RV (default).
+            For example, if you want to return the stellar RV relative to the barycenter, this is the planet mass. Note
+            that orbitize! assumes aop is that of the planet by default, so you would need to also add 180deg to
+            the aop to return stellar RV with the correct sign. If you want to return the planetary RV relative to 
+            the barycenter, this is the stellar mass. [Solar masses]. For planet mass ~ 0, mass_for_Kamp ~ M_tot, 
+            and function returns planetary RV (default).
         tau_ref_epoch (float, optional): reference date that tau is defined with respect to (default: 58849)
         tolerance (float, optional): absolute tolerance of iterative computation. Defaults to 1e-9.
         max_iter (int, optional): maximum number of iterations before switching. Defaults to 100.

orbitize/sampler.py

@@ -1001,10 +1001,10 @@ def run_sampler(
                     self._logl,
                     orbitize.priors.all_lnpriors,
                     ntemps=self.num_temps,
-                    threads=self.num_threads,
                     logpargs=[
                         self.priors,
                     ],
+                    pool=pool
                 )
             else:
                 sampler = emcee.EnsembleSampler(
@@ -1118,7 +1118,7 @@ def run_sampler(
                 self.results.save_results(output_filename)
 
             print("Run complete")
-        # Close pool
+
         if examine_chains:
             self.examine_chains()
 

orbitize/system.py

@@ -370,13 +370,14 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False):
                     Dec offsets from barycenter at each epoch.
 
                 vz (np.array of float): N_epochs x N_bodies x N_orbits array of
-                    radial velocities at each epoch.
+                    radial velocities at each epoch. RVs of the primary are
+                    relative to the barycenter, and RVs of secondary companions
+                    are relative to the primary.
 
                 brightness (np.array of float): N_epochs x N_bodies x N_orbits of 
                     photometric brightness predictions, assuming a Lambertian disk
                     reflection law, at each epoch. Normalized so that brightness=1
                     at maximum.
-
         """
 
         if epochs is None:
@@ -535,8 +536,12 @@ def compute_all_orbits(self, params_arr, epochs=None, comp_rebound=False):
                     vz0 = np.reshape(
                         vz_i * -(mass / m0), (n_epochs, n_orbits)
                     )  # calculating stellar velocity due to ith companion
-                    vz[:, 0, :] += vz0  # adding stellar velocity and gamma
-
+                    vz[:, 0, :] += vz0  # adding contribution from ith companion rv to stellar velocity
+
+            # Secondary RVs are assumed to be *relative* to the primary, so for all companions,
+            # we need to subtract the RV of the primary
+            vz[:,1:,:] -= vz[:, 0, :].reshape((n_epochs, 1, n_orbits))
+
             # if we are fitting for the mass of the planets, then they will perturb the star
             # add the perturbation on the star due to this planet on the relative astrometry of the planet that was measured
             # We are superimposing the Keplerian orbits, so we can add it linearly, scaled by the mass.

tests/test_multiplanet_rebound.py

@@ -1,6 +1,6 @@
 import os
 import numpy as np
-import astropy.table as table
+from astropy.table import Table, vstack
 import astropy.units as u
 import orbitize
 import orbitize.read_input as read_input
@@ -26,7 +26,7 @@
 
 def test_1planet():
     """
-    Sanity check that things agree for 1 planet case
+    Sanity check that things agree for 1 massive planet (really a star) case
     """
     # generate a planet orbit
     sma = 1
@@ -38,39 +38,63 @@ def test_1planet():
     plx = 1
     mtot = 1
     tau_ref_epoch = 0
-    mjup = u.Mjup.to(u.Msun)
-    mass_b = 12 * mjup
+    mass_b = 0.75
+    m0 = mtot - mass_b
 
     epochs = np.linspace(0, 300, 100) + tau_ref_epoch  # nearly the full period, MJD
 
-    ra_model, dec_model, vz_model = kepler.calc_orbit(
-        epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch
+    ra_model, dec_model, vz_st_model = kepler.calc_orbit(
+        epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch, mass_for_Kamp=mass_b
     )
+    ra_model, dec_model, vz_pl_model = kepler.calc_orbit(
+        epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=tau_ref_epoch, mass_for_Kamp=m0
+    )
+    vz_pl_model *= -1 # orbitize coordinate system (compute_model does this automatically)
 
-    # generate some fake measurements just to feed into system.py to test bookkeeping
-    t = table.Table(
+    # generate some fake measurements of the planet (relative astrom & relative secondary rv) 
+    # just to feed into system.py to test bookkeeping
+    t = Table(
         [
             epochs,
             np.ones(epochs.shape, dtype=int),
             ra_model,
             np.zeros(ra_model.shape),
             dec_model,
             np.zeros(dec_model.shape),
+            vz_pl_model - vz_st_model,
+            np.zeros(dec_model.shape),
         ],
-        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err"],
+        names=["epoch", "object", "raoff", "raoff_err", "decoff", "decoff_err", "rv", "rv_err"],
+    )
+    # add fake measurements of the planet (stellar rv)
+    t_rvs = Table(
+        [
+            epochs,
+            np.zeros(epochs.shape, dtype=int),
+            vz_st_model, np.zeros(vz_st_model.shape)
+        ],
+        names=["epoch", "object",  "rv", "rv_err"],
     )
+    t = vstack([t, t_rvs])
+
     filename = os.path.join(orbitize.DATADIR, "rebound_1planet.csv")
     t.write(filename, overwrite=True)
 
     # create the orbitize system and generate model predictions using the ground truth
-    astrom_dat = read_input.read_file(filename)
+    data = read_input.read_file(filename)
 
-    sys = system.System(1, astrom_dat, mtot, plx, tau_ref_epoch=tau_ref_epoch)
+    sys = system.System(1, data, mtot, plx, tau_ref_epoch=tau_ref_epoch, fit_secondary_mass=True)
 
-    params = np.array([sma, ecc, inc, aop, pan, tau, plx, mtot])
-    radec_orbitize, _ = sys.compute_model(params)
-    ra_orb = radec_orbitize[:, 0]
-    dec_orb = radec_orbitize[:, 1]
+    jit = 0
+    gamma = 0
+
+    params = np.array([sma, ecc, inc, aop, pan, tau, plx, gamma, jit, mass_b, m0])
+    modelpredict_orbitize, _ = sys.compute_model(params)
+
+    ra_orb = modelpredict_orbitize[:200:2, 0]
+    dec_orb = modelpredict_orbitize[:200:2, 1]
+    rv_pl_orb = modelpredict_orbitize[1:200:2,0]
+    rv_star_orb = modelpredict_orbitize[200:,0]
 
     # now project the orbit with rebound
     manom = basis.tau_to_manom(epochs[0], sma, mtot, tau, tau_ref_epoch)
@@ -95,24 +119,35 @@ def test_1planet():
     # integrate and measure star/planet separation
     ra_reb = []
     dec_reb = []
+    rv_star_reb = []
+    rv_pl_reb = []
 
     for t in epochs:
         sim.integrate(t / 365.25)
 
         ra_reb.append(-(ps[1].x - ps[0].x))  # ra is negative x
         dec_reb.append(ps[1].y - ps[0].y)
+        rv_star_reb.append(ps[0].vz)
+        rv_pl_reb.append(ps[1].vz - ps[0].vz)
 
     ra_reb = np.array(ra_reb)
     dec_reb = np.array(dec_reb)
+    rv_star_reb = np.array(rv_star_reb) * (u.au/u.yr).to(u.km/u.s)
+    rv_pl_reb = np.array(rv_pl_reb) * (u.au/u.yr).to(u.km/u.s)
+
 
     diff_ra = ra_reb - ra_orb / plx
     diff_dec = dec_reb - dec_orb / plx
+    diff_rv_st = rv_star_reb - rv_star_orb
+    diff_rv_pl = rv_pl_reb - rv_pl_orb
 
-    assert np.all(np.abs(diff_ra) < 1e-9)
-    assert np.all(np.abs(diff_dec) < 1e-9)
+    assert np.all(np.abs(diff_ra) < 1e-7)
+    assert np.all(np.abs(diff_dec) < 1e-7)
+    assert np.all(np.abs(diff_rv_st) < 1e-7)
+    assert np.all(np.abs(diff_rv_pl) < 1e-7)
 
     # clean up
-    # os.system("rm {}".format(filename))
+    os.system("rm {}".format(filename))
 
 
 def test_2planet_massive():
@@ -188,7 +223,7 @@ def test_2planet_massive():
 
     # generate some fake measurements of planet b, just to feed into system.py
     # to test bookkeeping
-    t = table.Table(
+    t = Table(
         [
             epochs,
             np.ones(epochs.shape, dtype=int),
@@ -298,7 +333,7 @@ def test_2planet_massive():
 
     # generate some fake measurements of planet c, just to feed into system.py to
     # test bookkeeping
-    t = table.Table(
+    t = Table(
         [
             epochs,
             np.ones(epochs.shape, dtype=int) * 2,
@@ -442,7 +477,7 @@ def test_2planet_massive_reverse_order():
 
     # generate some fake measurements of planet b, just to feed into system.py to test
     # bookkeeping
-    t = table.Table(
+    t = Table(
         [
             epochs,
             np.ones(epochs.shape, dtype=int) * 2,
@@ -599,7 +634,7 @@ def test_2planet_nomass():
 
     # generate some fake measurements of planet b, just to feed into system.py to
     # test bookkeeping
-    t = table.Table(
+    t = Table(
         [
             epochs,
             np.ones(epochs.shape, dtype=int),
@@ -698,6 +733,6 @@ def test_2planet_nomass():
 
 if __name__ == "__main__":
     test_1planet()
-    test_2planet_massive()
-    test_2planet_massive_reverse_order()
-    test_2planet_nomass()
+    # test_2planet_massive()
+    # test_2planet_massive_reverse_order()
+    # test_2planet_nomass()

tests/test_secondary_rvs.py

@@ -11,8 +11,9 @@ def test_secondary_rv_lnlike_calc():
     """
     Generates fake secondary RV data and asserts that
     the log(likelihood) of the true parameters is what we expect.
-    Also tests that the primary and secondary RV orbits are related by
-    -m/mtot
+    Also tests that the relative secondary RV output from compute_model
+    is the orbit of the secondary relative to the barycenter minus
+    the orbit of the primary relative to the barycenter.
     """
 
     # define an orbit & generate secondary RVs
@@ -23,20 +24,29 @@ def test_secondary_rv_lnlike_calc():
     Omega = 0
     tau = 0.3
     m0 = 1
-    m1 = 0.1
+    m1 = 0.7
     plx = 10
     orbitize_params_list = np.array([a, e, i, omega, Omega, tau, plx, m1, m0])
 
-    epochs = Time(np.linspace(2005, 2025, int(1e3)), format="decimalyear").mjd
+    epochs = Time(np.linspace(2005, 2025, int(1e1)), format="decimalyear").mjd
 
+    # compute RV of planet
     _, _, rv_p = calc_orbit(
         epochs, a, e, i, omega, Omega, tau, plx, m0 + m1, mass_for_Kamp=m0
     )
 
+    # compute RV of star
+    _, _, rv_s = calc_orbit(
+        epochs, a, e, i, omega+np.pi, Omega, tau, plx, m0 + m1, mass_for_Kamp=m1
+    )
+
+    # first check that the relationship between the RV datasets is what we expect
+    assert np.all(np.isclose(rv_s, -rv_p * m1/m0))
+
     data_file = DataFrame(columns=["epoch", "object", "rv", "rv_err"])
     data_file.epoch = epochs
     data_file.object = np.ones(len(epochs), dtype=int)
-    data_file.rv = rv_p
+    data_file.rv = rv_p - rv_s
     data_file.rv_err = np.ones(len(epochs)) * 0.01
 
     data_file.to_csv("tmp.csv", index=False)
@@ -62,7 +72,9 @@ def test_secondary_rv_lnlike_calc():
     rv0 = rv[:, 0]
     rv1 = rv[:, 1]
 
-    assert np.all(rv0 == -m1 / m0 * rv1)
+    # check that the output of compute_model for the secondary is the difference between the primary
+    # and secondary RV signals
+    assert np.all(np.isclose(rv1.flatten(), rv_p-rv_s))
 
 def test_read_input():
     """
@@ -74,6 +86,24 @@ def test_read_input():
     input_data['object'] = 1 # make sure all astrometry and RV is marked as of the secondary
     mySystem = system.System(1, input_data, 1, 1, fit_secondary_mass=False)
 
+def test_secondary_rvs_inst_specified():
+    """
+    Check that orbitize! sets up the System object correctly (i.e. doesn't fit for
+    gamma and jitter) even with multiple instruments specified for relative RVs.
+    """
+    input_data = read_input.read_file('{}/HD4747.csv'.format(DATADIR)) 
+    input_data['object'] = 1 # make sure all astrometry and RV is marked as of the secondary
+    input_data['inst'] = "test"
+    mySystem = system.System(1, input_data, 1, 1, fit_secondary_mass=False)
+
+    # check that inclusion of secondary RVs doesn't trigger gamma/jitter params to be created
+    assert not mySystem.fit_secondary_mass
+    for param in mySystem.param_idx.keys():
+        assert not param.startswith('gamma') and not param.startswith('jit')
+
+
 if __name__ == "__main__":
-    test_secondary_rv_lnlike_calc()
     test_read_input()
+
+    test_secondary_rv_lnlike_calc()
+    test_secondary_rvs_inst_specified()