diff --git a/docs/index.rst b/docs/index.rst index 47f8a6ea..354cf734 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -58,6 +58,7 @@ Changelog: - discuss MCMC autocorrelation in MCMC tutorial (@michaelkmpoon) - add time warning if OFTI doesn't accept an orbit in first 60 s (@michaelkmpoon) +- bugfix for rebound MCMC fits (issue #357; @sblunt) - implementation of Hipparcos-Gaia catalog of accelerations fitting! (@semaphoreP) - implementation of residual plotting method for orbit plots (@Saanikachoudhary and @semaphoreP) diff --git a/orbitize/nbody.py b/orbitize/nbody.py index 362b766b..1223ff08 100644 --- a/orbitize/nbody.py +++ b/orbitize/nbody.py @@ -2,118 +2,142 @@ import orbitize.basis as basis import rebound + def calc_orbit( - epochs, sma, ecc, inc, aop, pan, tau, plx, mtot, tau_ref_epoch=58849, - m_pl=None, output_star=False + epochs, + sma, + ecc, + inc, + aop, + pan, + tau, + plx, + mtot, + tau_ref_epoch=58849, + m_pl=None, + output_star=False, + integrator="ias15", ): """ Solves for position for a set of input orbital elements using rebound. Args: epochs (np.array): MJD times for which we want the positions of the planet - sma (np.array): semi-major axis of orbit [au] - ecc (np.array): eccentricity of the orbit [0,1] - inc (np.array): inclination [radians] - aop (np.array): argument of periastron [radians] - pan (np.array): longitude of the ascending node [radians] - tau (np.array): epoch of periastron passage in fraction of orbital period - past MJD=0 [0,1] - plx (np.array): parallax [mas] - mtot (np.array): total mass of the two-body orbit (M_* + M_planet) + sma (np.array): semi-major axis array of secondary bodies. For three planets, + this should look like: np.array([sma1, sma2, sma3]) [au] + ecc (np.array): eccentricity of the orbits (same format as sma) [0,1] + inc (np.array): inclinations (same format as sma) [radians] + aop (np.array): arguments of periastron (same format as sma) [radians] + pan (np.array): longitudes of the ascending node (same format as sma) [radians] + tau (np.array): epochs of periastron passage in fraction of orbital period + past MJD=0 (same format as sma) [0,1] + plx (float): parallax [mas] + mtot (float): total mass of the two-body orbit (M_* + M_planet) [Solar masses] tau_ref_epoch (float, optional): reference date that tau is defined with respect to - m_pl (np.array, optional): mass of the planets[units] + m_pl (np.array, optional): masss of the planets (same format as sma) [solar masses] output_star (bool): if True, also return the position of the star relative to the barycenter. + integrator (str): value to set for rebound.sim.integrator. Default "ias15" Returns: 3-tuple: - raoff (np.array): array-like (n_dates x n_orbs) of RA offsets between + raoff (np.array): array-like (n_dates x n_bodies x n_orbs) of RA offsets between the bodies (origin is at the other body) [mas] - deoff (np.array): array-like (n_dates x n_orbs) of Dec offsets between + deoff (np.array): array-like (n_dates x n_bodies x n_orbs) of Dec offsets between the bodies [mas] - - vz (np.array): array-like (n_dates x n_orbs) of radial velocity of - one of the bodies (see `mass_for_Kamp` description) [km/s] - .. Note:: - - return is in format [raoff[planet1, planet2,...,planetn], - deoff[planet1, planet2,...,planetn], vz[planet1, planet2,...,planetn] + vz (np.array): array-like (n_dates x n_bodies x n_orbs) of radial velocity of + one of the bodies (see `mass_for_Kamp` description) [km/s] """ - - sim = rebound.Simulation() #creating the simulation in Rebound - sim.units = ('AU','yr','Msun') #converting units for uniformity - sim.G = 39.476926408897626 #Using a more accurate value in order to minimize differences from prev. Kepler solver - ps = sim.particles #for easier calls - tx = len(epochs) #keeping track of how many time steps - te = epochs-epochs[0] #days + sim = rebound.Simulation() # creating the simulation in Rebound + sim.units = ("AU", "yr", "Msun") # converting units for uniformity + sim.G = 39.476926408897626 # Using a more accurate value in order to minimize differences from prev. Kepler solver + ps = sim.particles # for easier calls + + tx = len(epochs) # keeping track of how many time steps + te = epochs - epochs[0] # days - indv = len(sma) #number of planets orbiting the star - num_planets = np.arange(0,indv) #creates an array of indeces for each planet that exists + indv = len(sma) # number of planets orbiting the star + num_planets = np.arange( + 0, indv + ) # creates an array of indeces for each planet that exists - if m_pl is None: #if no planet masses are input, planet masses set ot zero and mass of star is equal to mtot - sim.add(m = mtot) + if ( + m_pl is None + ): # if no planet masses are input, planet masses set ot zero and mass of star is equal to mtot + sim.add(m=mtot) m_pl = np.zeros(len(sma)) m_star = mtot - else: #mass of star is always (mass of system)-(sum of planet masses) + else: # mass of star is always (mass of system)-(sum of planet masses) m_star = mtot - sum(m_pl) - sim.add(m = m_star) + sim.add(m=m_star) - #for each planet, create a body in the Rebound sim + # for each planet, create a body in the Rebound sim for i in num_planets: - #calculating mean anomaly - m_interior = m_star + sum(m_pl[0:i+1]) - mnm = basis.tau_to_manom(epochs[0], sma[i], m_interior, tau[i], tau_ref_epoch) - #adding each planet - sim.add(m = m_pl[i], a = sma[i], e = ecc[i], inc = inc[i], Omega = pan[i] + np.pi/2, omega =aop[i], M =mnm) + # calculating mean anomaly + m_interior = m_star + sum(m_pl[0 : i + 1]) + mnm = basis.tau_to_manom(epochs[0], sma[i], m_interior, tau[i], tau_ref_epoch) + # adding each planet + sim.add( + m=m_pl[i], + a=sma[i], + e=ecc[i], + inc=inc[i], + Omega=pan[i] + np.pi / 2, + omega=aop[i], + M=mnm, + ) sim.move_to_com() - sim.integrator = "ias15" - sim.dt = ps[1].P/100. #good rule of thumb: timestep should be at most 10% of the shortest orbital period + sim.integrator = integrator + sim.dt = ( + ps[1].P / 100.0 + ) # good rule of thumb: timestep should be at most 10% of the shortest orbital period if output_star: - ra_reb = np.zeros((tx, indv+1)) #numpy.zeros(number of [arrays], size of each array) - dec_reb = np.zeros((tx, indv+1)) - vz = np.zeros((tx, indv+1)) - for j,t in enumerate(te): - sim.integrate(t/365.25) - #for the star and each planet in each epoch denoted by j,t find the RA, Dec, and RV + ra_reb = np.zeros( + (tx, indv + 1) + ) # numpy.zeros(number of [arrays], size of each array) + dec_reb = np.zeros((tx, indv + 1)) + vz = np.zeros((tx, indv + 1)) + for j, t in enumerate(te): + sim.integrate(t / 365.25) + # for the star and each planet in each epoch denoted by j,t find the RA, Dec, and RV com = sim.com() - ra_reb[j,0] = -(ps[0].x-com.x)# ra is negative x - dec_reb[j,0] = ps[0].y-com.y - vz[j,0] = ps[0].vz + ra_reb[j, 0] = -(ps[0].x - com.x) # ra is negative x + dec_reb[j, 0] = ps[0].y - com.y + vz[j, 0] = ps[0].vz for i in num_planets: - ra_reb[j,i+1] = -(ps[int(i+1)].x - ps[0].x) # ra is negative x - dec_reb[j,i+1] = ps[int(i+1)].y - ps[0].y - vz[j,i+1] = ps[int(i+1)].vz + ra_reb[j, i + 1] = -(ps[int(i + 1)].x - ps[0].x) # ra is negative x + dec_reb[j, i + 1] = ps[int(i + 1)].y - ps[0].y + vz[j, i + 1] = ps[int(i + 1)].vz else: - ra_reb = np.zeros((tx, indv)) #numpy.zeros(number of [arrays], size of each array) + ra_reb = np.zeros( + (tx, indv) + ) # numpy.zeros(number of [arrays], size of each array) dec_reb = np.zeros((tx, indv)) vz = np.zeros((tx, indv)) - #integrate at each epoch - for j,t in enumerate(te): - sim.integrate(t/365.25) - #for each planet in each epoch denoted by j,t find the RA, Dec, and RV + # integrate at each epoch + for j, t in enumerate(te): + sim.integrate(t / 365.25) + # for each planet in each epoch denoted by j,t find the RA, Dec, and RV for i in num_planets: - ra_reb[j,i] = -(ps[int(i+1)].x - ps[0].x) # ra is negative x - dec_reb[j,i] = ps[int(i+1)].y - ps[0].y - vz[j,i] = ps[int(i+1)].vz - - - #adjusting for parallax - raoff = plx*ra_reb - deoff = plx*dec_reb - - #for formatting purposes - if len(sma)==1: - raoff = raoff.reshape(tx,) - deoff = deoff.reshape(tx,) - vz = vz.reshape(tx,) - return raoff, deoff, vz - else: - return raoff, deoff, vz \ No newline at end of file + ra_reb[j, i] = -(ps[int(i + 1)].x - ps[0].x) # ra is negative x + dec_reb[j, i] = ps[int(i + 1)].y - ps[0].y + vz[j, i] = ps[int(i + 1)].vz + + # adjusting for parallax + raoff = plx * ra_reb + deoff = plx * dec_reb + + # always assume we're using MCMC (i.e. n_orbits = 1) + raoff = raoff.reshape((tx, indv + 1, 1)) + deoff = deoff.reshape((tx, indv + 1, 1)) + vz = vz.reshape((tx, indv + 1, 1)) + + return raoff, deoff, vz diff --git a/tests/test_rebound.py b/tests/test_rebound.py index 82a2e91b..9f06afb8 100644 --- a/tests/test_rebound.py +++ b/tests/test_rebound.py @@ -1,5 +1,6 @@ from matplotlib import pyplot as plt from astropy.time import Time +from orbitize import sampler from orbitize.system import System from orbitize import DATADIR from orbitize.read_input import read_file @@ -125,8 +126,8 @@ def test_8799_rebound_vs_kepler(plotname=None): params_arr, epochs=epochs, comp_rebound=False ) - delta_ra = abs(rra - kra[:, :, 0]) - delta_de = abs(rde - kde[:, :, 0]) + delta_ra = abs(rra - kra) + delta_de = abs(rde - kde) yepochs = Time(epochs, format="mjd").decimalyear # check that the difference between these two solvers is smaller than @@ -179,8 +180,8 @@ def test_8799_rebound_vs_kepler(plotname=None): plt.plot(kra[:, 1:5, 0], kde[:, 1:5, 0], "indigo", label="Orbitize approx.") plt.plot(kra[-1, 1:5, 0], kde[-1, 1:5, 0], "o") - plt.plot(rra, rde, "r", label="Rebound", alpha=0.25) - plt.plot(rra[-1], rde[-1], "o", alpha=0.25) + plt.plot(rra[:, 1:5, 0], rde[:, 1:5, 0], "r", label="Rebound", alpha=0.25) + plt.plot(rra[-1, 1:5, 0], rde[-1, 1:5, 0], "o", alpha=0.25) plt.plot(0, 0, "*") plt.legend() @@ -199,5 +200,29 @@ def test_8799_rebound_vs_kepler(plotname=None): plt.savefig("{}_primaryorbittrack.png".format(plotname), dpi=250) +def test_rebound_mcmc(): + """ + Test that a rebound fit runs through one MCMC iteration successfully. + """ + + input_file = "{}/GJ504.csv".format(DATADIR) + data_table = read_file(input_file) + + my_sys = System( + num_secondary_bodies=1, + use_rebound=True, + fit_secondary_mass=True, + data_table=data_table, + stellar_or_system_mass=1.22, + mass_err=0, + plx=56.95, + plx_err=0.0, + ) + + my_mcmc_samp = sampler.MCMC(my_sys, num_temps=1, num_walkers=14, num_threads=1) + my_mcmc_samp.run_sampler(1, burn_steps=0) + + if __name__ == "__main__": - test_8799_rebound_vs_kepler(plotname="hr8799_diffs") + # test_8799_rebound_vs_kepler(plotname="hr8799_diffs") + test_rebound_mcmc()