From e4f3cf90f6459618cacd9af9e91f92220b843e30 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Mon, 12 Feb 2024 15:09:00 -0800 Subject: [PATCH 1/7] fix nbody indexing error & add unit test --- orbitize/nbody.py | 25 +++++++++---------------- tests/test_rebound.py | 34 +++++++++++++++++++++++++++++----- 2 files changed, 38 insertions(+), 21 deletions(-) diff --git a/orbitize/nbody.py b/orbitize/nbody.py index 362b766b..a203e456 100644 --- a/orbitize/nbody.py +++ b/orbitize/nbody.py @@ -30,19 +30,14 @@ def calc_orbit( 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 + 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] - - .. Note:: - - return is in format [raoff[planet1, planet2,...,planetn], - deoff[planet1, planet2,...,planetn], vz[planet1, planet2,...,planetn] """ sim = rebound.Simulation() #creating the simulation in Rebound @@ -109,11 +104,9 @@ def calc_orbit( 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 + # 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 \ No newline at end of file diff --git a/tests/test_rebound.py b/tests/test_rebound.py index 82a2e91b..7e7c8ac9 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,28 @@ def test_8799_rebound_vs_kepler(plotname=None): plt.savefig("{}_primaryorbittrack.png".format(plotname), dpi=250) +def test_rebound_mcmc(): + """ + Test that a 2-body rebound fit runs through one MCMC iteration successfully. + """ + + input_file = "{}/test_val_multi.csv".format(DATADIR) + data_table = read_file(input_file) + + my_sys = System(num_secondary_bodies=2, + use_rebound=True, + fit_secondary_mass=True, + data_table=data_table, + stellar_or_system_mass=1.0, mass_err=0, + plx=1.0, plx_err=0.0 + ) + + + my_mcmc_samp = sampler.MCMC(my_sys, num_temps=1, num_walkers=20, num_threads=1) + my_mcmc_samp.run_sampler(5, burn_steps=1) + + + if __name__ == "__main__": - test_8799_rebound_vs_kepler(plotname="hr8799_diffs") + # test_8799_rebound_vs_kepler(plotname="hr8799_diffs") + test_rebound_mcmc() From e661b5cfc897f41f033071849676f251622423a6 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Mon, 12 Feb 2024 15:21:46 -0800 Subject: [PATCH 2/7] more temps/walkers --- tests/test_rebound.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_rebound.py b/tests/test_rebound.py index 7e7c8ac9..406462a3 100644 --- a/tests/test_rebound.py +++ b/tests/test_rebound.py @@ -217,7 +217,7 @@ def test_rebound_mcmc(): ) - my_mcmc_samp = sampler.MCMC(my_sys, num_temps=1, num_walkers=20, num_threads=1) + my_mcmc_samp = sampler.MCMC(my_sys, num_temps=10, num_walkers=50, num_threads=1) my_mcmc_samp.run_sampler(5, burn_steps=1) From bc3e4cc39152679d4c9531afcb1b1ac454d2ee45 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Thu, 22 Feb 2024 19:23:12 -0800 Subject: [PATCH 3/7] lint --- orbitize/nbody.py | 138 ++++++++++++++++++++++++++++------------------ 1 file changed, 83 insertions(+), 55 deletions(-) diff --git a/orbitize/nbody.py b/orbitize/nbody.py index a203e456..c9f666c2 100644 --- a/orbitize/nbody.py +++ b/orbitize/nbody.py @@ -2,9 +2,20 @@ 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, ): """ Solves for position for a set of input orbital elements using rebound. @@ -16,10 +27,10 @@ def calc_orbit( 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 + 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) + mtot (np.array): 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 @@ -30,83 +41,100 @@ def calc_orbit( Returns: 3-tuple: - raoff (np.array): array-like (n_dates x n_bodies 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_bodies 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_bodies x n_orbs) of radial velocity of + + 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 = "ias15" + 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 - + 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 + # 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)) + vz = vz.reshape((tx, indv + 1, 1)) - return raoff, deoff, vz \ No newline at end of file + return raoff, deoff, vz From 56481f78a49590fd4720faf9c7c202a3c7a9be23 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Thu, 22 Feb 2024 19:31:16 -0800 Subject: [PATCH 4/7] set integrator nbody keyword --- orbitize/nbody.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/orbitize/nbody.py b/orbitize/nbody.py index c9f666c2..f92f5a22 100644 --- a/orbitize/nbody.py +++ b/orbitize/nbody.py @@ -16,6 +16,7 @@ def calc_orbit( tau_ref_epoch=58849, m_pl=None, output_star=False, + integrator="ias15", ): """ Solves for position for a set of input orbital elements using rebound. @@ -37,6 +38,7 @@ def calc_orbit( m_pl (np.array, optional): mass of the planets[units] 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: @@ -91,7 +93,7 @@ def calc_orbit( ) sim.move_to_com() - sim.integrator = "ias15" + 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 From f24b70f185da1cceb0d214aae05473de182ecc5e Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Thu, 22 Feb 2024 19:31:42 -0800 Subject: [PATCH 5/7] make rebound mcmc test faster --- tests/test_rebound.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/tests/test_rebound.py b/tests/test_rebound.py index 406462a3..9f06afb8 100644 --- a/tests/test_rebound.py +++ b/tests/test_rebound.py @@ -202,24 +202,25 @@ def test_8799_rebound_vs_kepler(plotname=None): def test_rebound_mcmc(): """ - Test that a 2-body rebound fit runs through one MCMC iteration successfully. + Test that a rebound fit runs through one MCMC iteration successfully. """ - input_file = "{}/test_val_multi.csv".format(DATADIR) + input_file = "{}/GJ504.csv".format(DATADIR) data_table = read_file(input_file) - my_sys = System(num_secondary_bodies=2, + my_sys = System( + num_secondary_bodies=1, use_rebound=True, fit_secondary_mass=True, data_table=data_table, - stellar_or_system_mass=1.0, mass_err=0, - plx=1.0, plx_err=0.0 + 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=10, num_walkers=50, num_threads=1) - my_mcmc_samp.run_sampler(5, burn_steps=1) - + 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__": From 1e1dd5ef73f9371d58674b58e2f4b773050088d3 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Fri, 5 Apr 2024 12:11:03 -0700 Subject: [PATCH 6/7] fix arg descriptions in nbody docstring --- orbitize/nbody.py | 21 +++++++++++---------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/orbitize/nbody.py b/orbitize/nbody.py index f92f5a22..1223ff08 100644 --- a/orbitize/nbody.py +++ b/orbitize/nbody.py @@ -23,19 +23,20 @@ def calc_orbit( 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" From 4da8772a6a2747c72661c0fcd3e1c4e6ed314407 Mon Sep 17 00:00:00 2001 From: Sarah Blunt Date: Wed, 10 Apr 2024 12:51:11 -0700 Subject: [PATCH 7/7] add issue 357 fix to changelog --- docs/index.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/index.rst b/docs/index.rst index 1ca2437e..02df34e6 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) **2.2.2 (2023-06-30)**