From 0b463938ef407f470854bc9ff9a5532fa03dfa7e Mon Sep 17 00:00:00 2001 From: dyahalomi Date: Thu, 9 Feb 2023 15:51:17 -0700 Subject: [PATCH 1/6] Add files via upload --- docs/tutorials/astrometry_rv_case_study.ipynb | 2029 +++++++++++++++++ 1 file changed, 2029 insertions(+) create mode 100644 docs/tutorials/astrometry_rv_case_study.ipynb diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb new file mode 100644 index 00000000..4e43e340 --- /dev/null +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -0,0 +1,2029 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "fae3c154", + "metadata": {}, + "source": [ + "# Joint Astrometric and Radial Velocity Fitting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4abd3961", + "metadata": {}, + "outputs": [], + "source": [ + "import exoplanet as xo\n", + "import numpy as np\n", + "\n", + "\n", + "xo.utils.docs_setup()\n", + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "9a87ff47", + "metadata": {}, + "source": [ + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", + "\n", + "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", + "\n", + "\n", + "RV observations have the downside that they can only determine the minimum mass of an exoplanet as RV observations have only 1 spacial dimension (towards and away from the observer) and 1 timelike dimension. Therefore, there exists a mass-inclination degeneracy. Luckily, we can hope to break this mass-inclination degeneracy with astrometric observations.\n", + "\n", + "\n", + "\n", + "Gaia will release 100-200 individual epoch astrometry observations for these THE stars (as they will be bright and nearby) in data release 4. These observations will have a 1$\\sigma$ measurement error of approximately 60 $\\mu$as (micro-arcseconds). You can read more about Gaia on their website: https://sci.esa.int/web/gaia or in the many papers written about Gaia and its first 3 data releases.\n", + "\n", + "Additionally, The Nancy Grace Roman Space Telescope, which is set to launch in 2027 will be capable of precision astrometry (approximately 5-20 $\\mu$as) using either the diffraction spike method or the spacial scanning method. In addition to having a higher expected precision per observation than Gaia, Roman astrometry will be very helpful to the fits by extending the astrometric baseline up to ~25 years. To read more about Roman, see their website: https://roman.gsfc.nasa.gov.\n", + "\n", + "\n", + "In what follows, we model in right ascension and declination space (rather than the standard position angle and seperation used in astrometry modeling). We made this decision as position angle doesn't have an independent uncertainty since it depends on the amplitude of the separation, and so we thought it would be a better modeling choice to simulate and model the data in right ascension and declination space." + ] + }, + { + "cell_type": "markdown", + "id": "93ceb2a9", + "metadata": {}, + "source": [ + "## (Simulated) Data\n", + "\n", + "\n", + "First, we need to simulate our dataset. \n", + "\n", + "\n", + "### Let's define 2 helpful functions in order to simulate our data...." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2325bae1", + "metadata": {}, + "outputs": [], + "source": [ + "def create_orbit(n_planets, orbit_params):\n", + " \n", + " '''\n", + " inputs:\n", + " - n_planets = number of planets\n", + " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", + " - period = period in days\n", + " - ecc = eccentricity \n", + " - t_periastron = The epoch of a reference periastron passage in days.\n", + " - omega = argument of periapsis in radans\n", + " - Omega = longitude of ascending node \n", + " - incl = inclination in radians\n", + " - m_planet = mass of planet in Earth masses \n", + " '''\n", + " \n", + "\n", + " \n", + " period = []\n", + " ecc = []\n", + " Tper = [] \n", + " omega = [] \n", + " Omega = [] \n", + " incl = []\n", + " m_planet = []\n", + " for ii in range(0, n_planets):\n", + " period.append(orbit_params[ii][0])\n", + " ecc.append(orbit_params[ii][1])\n", + " Tper.append(orbit_params[ii][2])\n", + " omega.append(orbit_params[ii][3])\n", + " Omega.append(orbit_params[ii][4])\n", + " incl.append(orbit_params[ii][5])\n", + " m_planet.append(orbit_params[ii][6])\n", + "\n", + " orbit = xo.orbits.KeplerianOrbit(\n", + " period = period,\n", + " ecc = ecc, \n", + " t_periastron = Tper, \n", + " omega = omega, \n", + " Omega = Omega, \n", + " incl = incl,\n", + " m_planet = m_planet) \n", + " \n", + " \n", + " return orbit" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a70101ef", + "metadata": {}, + "outputs": [], + "source": [ + "np.random.seed(979) \n", + "\n", + "def simulate_data(\n", + " n_planets, \n", + " sigma_rv, \n", + " sigma_ra,\n", + " sigma_dec,\n", + " plx,\n", + " orbit_params,\n", + " times_observed_rv = None,\n", + " times_observed_astrometry = None,\n", + " t_dur_rv = None, \n", + " n_obs_rv = None,\n", + " t_dur_astrometry = None, \n", + " n_obs_astrometry = None \n", + " ):\n", + "\n", + " '''\n", + " inputs:\n", + " - orbit = keplerian orbit generated by exoplanet code\n", + " - n_planets = number of planets\n", + " - sigma_rv = standard deviation for RV simulated gaussian noise in as\n", + " - sigma_theta = standard deviation for position angle simulated gaussian noise in as\n", + " - sigma_rho = standard deviation for seperation simulated gaussian noise in as\n", + " - plx = parallax of system\n", + " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", + " - period = period in days\n", + " - ecc = eccentricity \n", + " - t_periastron = The epoch of a reference periastron passage in days.\n", + " - omega = argument of periapsis in radans\n", + " - Omega = longitude of ascending node \n", + " - incl = inclination in radians\n", + " - m_planet = mass of planet in Earth masses \n", + " - times_observed_rv = (default=None) the observed times for RV observations\n", + " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", + " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", + " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", + " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", + " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", + " \n", + " \n", + " '''\n", + " \n", + " if n_planets == 1:\n", + " orbit_params = [orbit_params]\n", + " \n", + " \n", + " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", + "\n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined\n", + " else:\n", + " start_time_rv = np.inf\n", + " for ii in range(0, n_planets):\n", + " if orbit_params[ii][2] < start_time_rv:\n", + " start_time_rv = orbit_params[ii][2]\n", + "\n", + "\n", + " \n", + " #define RV observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", + " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", + " \n", + "\n", + "\n", + " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", + " \n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined \n", + " else:\n", + " start_time_astrometry = np.inf\n", + " for ii in range(0, n_planets):\n", + " if orbit_params[ii][2] < start_time_astrometry:\n", + " start_time_astrometry = orbit_params[ii][2] \n", + "\n", + "\n", + " #define astrometry observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", + " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", + "\n", + "\n", + " \n", + " orbit = create_orbit(n_planets, orbit_params)\n", + " \n", + " rv = orbit.get_radial_velocity(times_rv)\n", + " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", + " \n", + " rv_orbit = rv.eval()\n", + "\n", + " if n_planets > 1:\n", + " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", + " else:\n", + " rv_orbit_sum = rv_orbit\n", + " \n", + "\n", + " rv_observed = rv_observed.eval()\n", + " if n_planets > 1:\n", + " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", + " else:\n", + " rv_observed_sum = rv_observed\n", + "\n", + "\n", + " #-----------\n", + " #-----------\n", + "\n", + " # determine and print the star position at desired times\n", + " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", + "\n", + "\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x,y,z = pos\n", + "\n", + "\n", + " # calculate rho and theta\n", + " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", + " \n", + " rho, theta = rho.eval(), theta.eval()\n", + " \n", + " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", + "\n", + "\n", + " rho_orbit = rho\n", + " \n", + " if n_planets > 1:\n", + " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", + " else:\n", + " rho_orbit_sum = rho_orbit\n", + "\n", + " theta_orbit = theta\n", + " \n", + " if n_planets > 1:\n", + " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", + " else:\n", + " theta_orbit_sum = theta_orbit\n", + " \n", + " \n", + " # when summing over theta, position angle, we have to careful because position\n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", + " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", + " # in the correct range. Be careful though if modeling more than 2 planets, this\n", + " # doesn't completely solve the problem!\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", + "\n", + "\n", + " # determine and print the star position at desired times\n", + " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", + "\n", + "\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x_obs,y_obs,z_obs = pos_observed\n", + "\n", + "\n", + " # calculate rho and theta\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", + " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", + " \n", + " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", + "\n", + "\n", + " \n", + " if n_planets > 1:\n", + " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", + " else:\n", + " rho_observed_sum = rho_observed\n", + "\n", + " \n", + " if n_planets > 1:\n", + " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", + " else:\n", + " theta_observed_sum = theta_observed\n", + " \n", + " \n", + " #######\n", + " #######\n", + " #######\n", + "\n", + " # convert rho and theta into ra and dec\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", + "\n", + " if n_planets > 1:\n", + " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", + " else:\n", + " ra_orbit_sum = ra_orbit\n", + "\n", + " \n", + " if n_planets > 1:\n", + " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", + " else:\n", + " dec_orbit_sum = dec_orbit\n", + "\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", + "\n", + "\n", + "\n", + "\n", + " \n", + " #----------- \n", + " #----------- \n", + " #----------- simulate data\n", + " #----------- \n", + " #----------- \n", + "\n", + "\n", + " if n_planets > 1:\n", + " rv_sim = rv_observed\n", + " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", + " else:\n", + " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", + " rv_sim_sum = rv_sim\n", + "\n", + " #-----------\n", + " #-----------\n", + "\n", + "\n", + " if n_planets > 1:\n", + " ra_sim = ra_observed \n", + " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", + " \n", + " else:\n", + " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", + " ra_sim_sum = ra_sim\n", + " \n", + " \n", + " \n", + " if n_planets > 1:\n", + " dec_sim = dec_observed \n", + " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " \n", + " else:\n", + " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " dec_sim_sum = dec_sim\n", + "\n", + " \n", + " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", + " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", + " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", + " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", + " \n", + " \n", + " return [times, rv_results, ra_results, dec_results]" + ] + }, + { + "cell_type": "markdown", + "id": "9d5d78fc", + "metadata": {}, + "source": [ + "### Let's now define our orbital parameters:\n", + "\n", + "We will be simulating (and then modeling) the signal from Earth and Jupiter orbiting around the Sun, observed 10 parsecs (parallax = 0.1 as) away, with an inclination of 45 degrees. In total, our orbital parameters are:\n", + "\n", + "|planet | period (days) | ecc | Tper | omega (rad) | Omega (rad) | inclination (rad) | mass (solar masses) |\n", + "| --- | --- | --- | --- | --- | --- | --- | --- |\n", + "|Earth | 365.0 | 0.0167 | 100.0 | 1.7959438003021653 | 0.0 | $\\pi$/4 | 3e-06 |\n", + "|Jupiter | 4327.0 | 0.0484 | 500.0 | -1.4957471689591395 | 1.7523105690023069 | $\\pi$/4 + 0.022759 | 0.000955 |" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "68e45e39", + "metadata": {}, + "outputs": [], + "source": [ + "import numpy as np\n", + "\n", + "parallax = 0.1 #as \n", + "n_planets=2\n", + "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", + " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" + ] + }, + { + "cell_type": "markdown", + "id": "1b8b3007", + "metadata": {}, + "source": [ + "### And now we define our observing parameters:\n", + "\n", + "Terra Hunting RVs: Nightly observations for 10 years, half a year each year (as the other half of the time the star won't be visible from the ground). White noise errors = 1$\\sigma$ measurement error = 0.3 m/s.\n", + "\n", + "Gaia Astrometry: To be conservative, we will assume we have 100 Gaia observations spread out evenly over the 10 years since the initial launch in 2013. White noise errors = 1$\\sigma$ measurement error = 60 $\\mu$as.\n", + "\n", + "Roman Astrometry: Here we need to assume both an observing strategy (cadence and duration) and a white noise error. In this example we will assume that Roman observes for 10 years, 4 times a year, with a white noise errors = 1$\\sigma$ measurement error = 5 $\\mu$as. We tried other values here, as presented in Yahalomi et al. 2023.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "55938572", + "metadata": {}, + "outputs": [], + "source": [ + "from aesara_theano_fallback import aesara as theano\n", + "import theano.tensor as tt\n", + "\n", + "\n", + "#add gaia observing times\n", + "times_observed_astrometry_gaia = []\n", + "t_0 = 0\n", + "for ii in range(t_0, t_0+3600):\n", + " if ii % (3600/100) == 0:\n", + " times_observed_astrometry_gaia.append(ii)\n", + "\n", + "\n", + "\n", + "#add THE observing times\n", + "times_observed_rv = []\n", + "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "add_data = True\n", + "for ii in range(t_0, t_0+3600):\n", + "\n", + " if ii % 180 == 0:\n", + " if add_data:\n", + " add_data = False\n", + " else:\n", + " add_data = True\n", + "\n", + " if add_data:\n", + " times_observed_rv.append(ii)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "sigma_rv = 0.3\n", + "\n", + "sigma_ra_gaia = 6e-5 #60 micro-as\n", + "sigma_dec_gaia = 6e-5 #60 micro-as\n", + "\n", + "\n", + "\n", + "times, rv_results, ra_results, dec_results = simulate_data(\n", + " n_planets, \n", + " sigma_rv, \n", + " sigma_ra_gaia,\n", + " sigma_dec_gaia,\n", + " parallax,\n", + " planet_params,\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_gaia\n", + " )\n", + "\n", + "\n", + "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", + "\n", + "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", + "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "# we assume that the first roman observation\n", + "# will occur ~5 years after the last Gaia observation\n", + "t_1 = times_observed_astrometry_gaia[-1]+1800\n", + "\n", + "# add roman observing times\n", + "times_observed_astrometry_roman = []\n", + "for ii in range(t_1, t_1+(10*365)):\n", + " if ii % 90 == 0:\n", + " times_observed_astrometry_roman.append(ii) \n", + "\n", + "\n", + "sigma_ra_roman = 5e-6 #5 micro-as\n", + "sigma_dec_roman = 5e-6 #5 micro-as\n", + "\n", + "\n", + "\n", + "times, rv_results, ra_results, dec_results = simulate_data(\n", + " n_planets, \n", + " sigma_rv, \n", + " sigma_ra_roman,\n", + " sigma_dec_roman,\n", + " parallax,\n", + " planet_params,\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_roman\n", + " )\n", + "\n", + "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", + "\n", + "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", + "\n", + "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", + "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", + "ra_sim = np.append(ra_sim, ra_results[2], axis=0)\n", + "ra_sim_sum = np.append(ra_sim_sum, ra_results[3], axis=0)\n", + "\n", + "dec_orbit = np.append(dec_orbit, dec_results[0], axis=0)\n", + "dec_orbit_sum = np.append(dec_orbit_sum, dec_results[1], axis=0)\n", + "dec_sim = np.append(dec_sim, dec_results[2], axis=0)\n", + "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", + "\n", + "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" + ] + }, + { + "cell_type": "markdown", + "id": "d4af3858", + "metadata": {}, + "source": [ + "## Let's write some functions to plot our simulated data:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "21588382", + "metadata": {}, + "outputs": [], + "source": [ + "import matplotlib.pyplot as plt\n", + "import matplotlib\n", + "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", + "matplotlib.rc('text', usetex=True)\n", + "\n", + "def plot_rv_signal(\n", + " n_planets, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv):\n", + " \n", + " \n", + " \n", + " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", + " fig.suptitle(\"RV Signal\", fontsize = 45)\n", + " ax1 = ax\n", + " \n", + " \n", + " \n", + "\n", + " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", + "\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", + " ax1.legend(fontsize = 27, loc = 2)\n", + "\n", + " fig.tight_layout()\n", + " fig.show()\n", + " \n", + " return None" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33072448", + "metadata": {}, + "outputs": [], + "source": [ + "def plot_astrometry_signal(\n", + " n_planets, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", + " dec_sim_sum,\n", + " times_astrometry,\n", + " times_observed_astrometry):\n", + " \n", + " \n", + "\n", + " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", + " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", + " ax2 = ax[0]\n", + " ax3 = ax[1]\n", + " \n", + "\n", + "\n", + " \n", + " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " for tick in ax3.get_xticklabels():\n", + " tick.set_rotation(30)\n", + " ax2.legend(fontsize = 27, loc = 2)\n", + " \n", + " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax3.legend(fontsize = 27, loc = 2)\n", + "\n", + "\n", + "\n", + " fig.tight_layout()\n", + " fig.show()\n", + "\n", + "\n", + "\n", + " #######\n", + " #######\n", + " #######\n", + "\n", + "\n", + " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", + "\n", + " # plot simulated data\n", + " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", + "\n", + " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", + " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", + " ax.invert_xaxis()\n", + " ax.plot(0, 0, \"k*\")\n", + " ax.set_aspect(\"equal\", \"datalim\")\n", + " ax.set_title(\"initial orbit\")\n", + " ax.legend()\n", + " fig.show()\n", + " \n", + " return None" + ] + }, + { + "cell_type": "markdown", + "id": "1b967e32", + "metadata": {}, + "source": [ + "## And now, let's plot" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "82e18e64", + "metadata": {}, + "outputs": [], + "source": [ + "plot_rv_signal(\n", + " 2, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv)\n", + " \n", + " \n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "623b9c2d", + "metadata": {}, + "outputs": [], + "source": [ + "plot_astrometry_signal(\n", + " 2, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", + " dec_sim_sum,\n", + " times_astrometry,\n", + " times_observed_astrometry)" + ] + }, + { + "cell_type": "markdown", + "id": "99110cc0", + "metadata": {}, + "source": [ + "## Ok, now we have a simulated dataset -- let's model it!\n", + "\n", + "\n", + "Let's start with an RV only model -- this will give us an initial guess for many of the parameters in the dataset and speed up the joint model later on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a671d599", + "metadata": {}, + "outputs": [], + "source": [ + "def minimize_rv(periods, Ks, x_rv, y_rv, y_rv_err):\n", + " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", + " print(\"minimizing RV only model solutions pre-MCMC\")\n", + " print(\"------------\")\n", + "\n", + " with pm.Model() as model:\n", + " \n", + " \n", + " #log normal prior on period around estimates\n", + " logP = pm.Uniform(\n", + " \"logP\",\n", + " lower=0,\n", + " upper=11,\n", + " shape=2,\n", + " testval=np.log(periods),\n", + " )\n", + " \n", + "\n", + "\n", + " P = pm.Deterministic(\"P\", tt.exp(logP))\n", + "\n", + "\n", + " ## wide uniform prior on t_periastron\n", + " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", + " \n", + " \n", + " # Wide normal prior for semi-amplitude\n", + " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", + " \n", + " K = pm.Deterministic(\"K\", tt.exp(logK))\n", + " \n", + " \n", + " # Eccentricity & argument of periasteron\n", + " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", + "\n", + " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", + "\n", + "\n", + "\n", + " \n", + " # Jitter & a quadratic RV trend\n", + " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + "\n", + " # Then we define the orbit\n", + " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", + "\n", + " # And a function for computing the full RV model\n", + " def get_rv_model(t, name=\"\"):\n", + " # First the RVs induced by the planets\n", + " vrad = orbit.get_radial_velocity(t, K=K)\n", + " pm.Deterministic(\"vrad\" + name, vrad)\n", + "\n", + " # Sum over planets and add the background to get the full model\n", + " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", + "\n", + " \n", + "\n", + " # Define the RVs at the observed times\n", + " rv_model = get_rv_model(x_rv)\n", + "\n", + " # Also define the model on a fine grid as computed above (for plotting)\n", + " rv_model_pred = get_rv_model(t_rv, name=\"_pred\")\n", + "\n", + " # Finally add in the observation model. This next line adds a new contribution\n", + " # to the log probability of the PyMC3 model\n", + " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " err = y_rv_err\n", + " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", + "\n", + "\n", + " map_soln = model.test_point\n", + " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", + " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", + " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", + " map_soln = pmx.optimize(start=map_soln)\n", + "\n", + "\n", + " #return the max a-posteriori solution\n", + " return map_soln" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4d6faa40", + "metadata": {}, + "outputs": [], + "source": [ + "import pymc3 as pm\n", + "import pymc3_ext as pmx\n", + "\n", + "np.random.seed(979) \n", + "\n", + "\n", + "#rename variables in more consistent way for modeling\n", + "x_rv = np.array(times_observed_rv)\n", + "y_rv = rv_sim_sum\n", + "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", + "\n", + "x_astrometry = np.array(times_observed_astrometry)\n", + "ra_data = ra_sim_sum\n", + "dec_data = dec_sim_sum\n", + "\n", + "\n", + "ra_err = np.concatenate((ra_gaia_err, ra_roman_err))\n", + "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", + "\n", + "\n", + " \n", + " \n", + "# make a fine grid that spans the observation window for plotting purposes\n", + "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", + "\n", + "# for predicted orbits\n", + "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "\n", + "\n", + "#minimize on RV data\n", + "periods_real = [4327.0, 365.0]\n", + "periods_guess = []\n", + "\n", + "#let's assume that THE can recover the period within 10% of the true value\n", + "for per in periods_real:\n", + " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", + "\n", + "print('periods used as starting point = ', periods_guess)\n", + "periods_guess = np.array(periods_guess)\n", + "\n", + "\n", + "Ks_guess = xo.estimate_semi_amplitude(periods_guess, x_rv, y_rv, y_rv_err)\n", + "\n", + "rv_map_soln = minimize_rv(periods_guess, Ks_guess, x_rv, y_rv, y_rv_err)" + ] + }, + { + "cell_type": "markdown", + "id": "efc4dcc2", + "metadata": {}, + "source": [ + "## Now let's define out joint astrometry and radial velocity model:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "759b3536", + "metadata": {}, + "outputs": [], + "source": [ + "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", + " m_sun = 333030 #earth masses\n", + " \n", + " P_RV = np.array(rv_map_soln['P'])\n", + " K_RV = np.array(rv_map_soln['K'])\n", + " tperi_RV = np.array(rv_map_soln['tperi'])\n", + " ecc_RV = np.array(rv_map_soln['ecc'])\n", + " omega_RV = np.array(rv_map_soln['omega'])\n", + " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", + " phase_RV = determine_phase(P_RV, tperi_RV)\n", + " \n", + " \n", + " # make a fine grid that spans the observation window for plotting purposes\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", + "\n", + " # for predicted orbits\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "\n", + "\n", + " print(\"RV Solutions\")\n", + " print(\"------------\")\n", + " print(\"P: \", P_RV)\n", + " print(\"K: \", K_RV)\n", + " print(\"T_peri: \", tperi_RV)\n", + " print(\"eccentricity: \", ecc_RV)\n", + " print(\"omega: \", omega_RV)\n", + "\n", + " print('')\n", + " print(\"minimizing joint model solutions pre-MCMC\")\n", + " print(\"------------\")\n", + "\n", + " # for predicted orbits\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "\n", + " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", + " model, map_soln, logp = [], [], []\n", + " for inc in inc_test_vals:\n", + " mass_test_vals = min_masses_RV/np.sin(inc)\n", + " print('')\n", + " print(\"trying inclination = \" + str(np.degrees(inc)))\n", + " print(\"mass test val = \" + str(mass_test_vals))\n", + " print(\"------------\")\n", + " \n", + " \n", + "\n", + "\n", + "\n", + " def get_model():\n", + " with pm.Model() as model:\n", + "\n", + "\n", + " # Below we will run a version of this model where a measurement of parallax is provided\n", + " # The measurement is in milliarcsec\n", + " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", + " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", + " )\n", + " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", + "\n", + "\n", + " # We expect the period to be around that found from just the RVs,\n", + " # so we'll set a broad prior on logP\n", + " \n", + " logP = pm.Uniform(\n", + " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", + " )\n", + " P = pm.Deterministic(\"P\", tt.exp(logP))\n", + " \n", + " # Eccentricity & argument of periasteron\n", + " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", + " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", + " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", + "\n", + " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", + "\n", + " \n", + " \n", + "\n", + " # Omegas are co-dependent, so sample them with variables Omega_plus\n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", + " # Omega_minus is (Omega_0 - Omega_1)/2\n", + " \n", + " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", + " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", + " \n", + " \n", + " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", + " (Omega_plus - Omega_minus)] )\n", + " \n", + "\n", + " Omega = pm.Deterministic(\"Omega\", Omega) \n", + " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", + " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", + " \n", + "\n", + "\n", + " \n", + " # For these orbits, it can also be better to fit for a phase angle\n", + " # (relative to a reference time) instead of the time of periasteron\n", + " # passage directly\n", + " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", + " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", + " \n", + "\n", + " \n", + " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", + " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", + " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", + " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", + "\n", + "\n", + " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", + " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", + " \n", + "\n", + "\n", + " \n", + " \n", + " # Set up the orbit\n", + " orbit = xo.orbits.KeplerianOrbit(\n", + " t_periastron=tperi,\n", + " period=P,\n", + " incl=incl,\n", + " ecc=ecc,\n", + " omega=omega,\n", + " Omega=Omega,\n", + " m_planet = m_planet_fit,\n", + " plx=plx\n", + " )\n", + "\n", + "\n", + " \n", + " \n", + " # Add a function for computing the full astrometry model\n", + " def get_astrometry_model(t, name=\"\"):\n", + " # First the astrometry induced by the planets\n", + "\n", + " # determine and print the star position at desired times\n", + " pos = orbit.get_star_position(t, plx)\n", + "\n", + " x,y,z = pos\n", + "\n", + "\n", + " # calculate rho and theta\n", + " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", + " \n", + " \n", + " #rhos, thetas = get_star_relative_angles(t, plx)\n", + " \n", + " \n", + " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", + " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", + " \n", + " # Sum over planets to get the full model\n", + " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", + " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", + " \n", + "\n", + " \n", + " return dec_model, ra_model\n", + "\n", + " \n", + " # Define the astrometry model at the observed times\n", + " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", + "\n", + " # Also define the model on a fine grid as computed above (for plotting)\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", + "\n", + " \n", + "\n", + "\n", + " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " dec_tot_err = dec_err\n", + " ra_tot_err = ra_err\n", + "\n", + "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", + " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", + " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", + "\n", + "\n", + " \n", + " \n", + " # ADD RV MODEL\n", + "\n", + "\n", + " # And a function for computing the full RV model\n", + " def get_rv_model(t, name=\"\"):\n", + " # First the RVs induced by the planets\n", + " vrad = orbit.get_radial_velocity(t)\n", + " pm.Deterministic(\"vrad\" + name, vrad)\n", + "\n", + " # Sum over planets to get the full model\n", + " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", + "\n", + " # Define the RVs at the observed times\n", + " rv_model = get_rv_model(x_rv)\n", + "\n", + " # Also define the model on a fine grid as computed above (for plotting)\n", + " rv_model_pred = get_rv_model(t_rv, name=\"_pred\")\n", + "\n", + " # Finally add in the observation model. This next line adds a new contribution\n", + " # to the log probability of the PyMC3 model\n", + " rv_err = y_rv_err\n", + " pm.Normal(\"obs_RV\", mu=rv_model, sd=rv_err, observed=y_rv)\n", + "\n", + " # Optimize to find the initial parameters\n", + " map_soln = model.test_point\n", + " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", + " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + "\n", + " map_soln = pmx.optimize(map_soln)\n", + "\n", + "\n", + " return model, map_soln\n", + "\n", + "\n", + "\n", + " a_model, a_map_soln = get_model()\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", + " print('log likelihood = ' + str(a_logp))\n", + "\n", + " model.append(a_model)\n", + " map_soln.append(a_map_soln)\n", + " logp.append(a_logp)\n", + "\n", + " \n", + " best_index = 0\n", + " for index in range(0, len(model)):\n", + " if logp[index] >= logp[best_index]:\n", + " best_index = index\n", + "\n", + " the_model = model[best_index]\n", + " the_map_soln = map_soln[best_index]\n", + " the_logp = logp[best_index]\n", + "\n", + "\n", + "\n", + " return the_model, the_map_soln, the_logp\n", + "\n", + "\n", + "\n", + "def determine_phase(P, t_periastron):\n", + " phase = (2 * np.pi * t_periastron) / P\n", + " return phase\n", + "\n", + "\n", + "\n", + "def model_both(model, map_soln, tune_steps, draw_steps):\n", + " print('Joint RV + Astometry Minimization Solutions:')\n", + " print(\"------------\")\n", + "\n", + " print('m_planet: ' + str(map_soln['m_planet']))\n", + " print('P: ' + str(map_soln['P']))\n", + " print('incl: ' + str(map_soln['incl']))\n", + " print('Omega: ' + str(map_soln['Omega']))\n", + " print('tperi: ' + str(map_soln['tperi']))\n", + " print('ecc: ' + str(map_soln['ecc']))\n", + " print('omega: ' + str(map_soln['omega']))\n", + " print('plx: ' + str(map_soln['plx']))\n", + "\n", + "\n", + " with model:\n", + " trace = pmx.sample(\n", + " tune=tune_steps,\n", + " draws=draw_steps,\n", + " start=map_soln,\n", + " cores=2,\n", + " chains=2,\n", + " target_accept=0.99,\n", + " return_inferencedata=True,\n", + " )\n", + "\n", + " return trace" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "041c8f43", + "metadata": {}, + "outputs": [], + "source": [ + "from astropy.constants import M_sun\n", + "from astropy import constants\n", + "\n", + "################\n", + "################\n", + "#minimize on joint model\n", + "joint_model, joint_map_soln, joint_logp = minimize_both(\n", + " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", + " ra_data, ra_err, dec_data, dec_err, parallax\n", + ")\n", + "\n", + "\n", + "\n", + "\n", + "################\n", + "################\n", + "#run full MCMC\n", + "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" + ] + }, + { + "cell_type": "markdown", + "id": "72fa59d9", + "metadata": {}, + "source": [ + "## Define some plotting functions that will help us check convergence and the covariance of the modelling parameters:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4d24a32e", + "metadata": {}, + "outputs": [], + "source": [ + "matplotlib.rc('text', usetex=False)\n", + "\n", + "def make_plots(trace, orbit_params, n_planets):\n", + " m_sun = 333030 #earth masses\n", + "\n", + " [orbit_params_earth, orbit_params_jup] = orbit_params\n", + "\n", + "\n", + " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", + " inclination_earth, m_earth] = orbit_params_earth\n", + "\n", + "\n", + " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", + " inclination_jup, m_jup] = orbit_params_jup\n", + "\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", + " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", + " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", + "\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", + " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # plot the table summarizing MCMC results\n", + " az.summary(\n", + " trace,\n", + " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", + " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", + " )\n", + "\n", + "\n", + " plt.show()\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # plot the corner plots\n", + " _ = corner.corner(\n", + " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", + " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", + " truths = truths, truth_color = \"#03003a\"\n", + " )\n", + "\n", + " plt.show()\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # plot the chains\n", + " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", + " for ii in range(0, len(parameters)):\n", + " plot_truth = False\n", + " param = parameters[ii]\n", + " \n", + " true_vals_earth = truth_chain_plot[2*ii]\n", + " true_vals_jup = truth_chain_plot[2*ii+1]\n", + " plot_truth = True\n", + " \n", + " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", + " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", + " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", + " \n", + " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", + " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", + " \n", + " \n", + " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", + " \n", + " \n", + " ax[0].plot(nstep, planet1_chain1)\n", + " ax[0].plot(nstep, planet1_chain2)\n", + " \n", + " if plot_truth:\n", + " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", + " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", + " ax[0].legend(fontsize = 18)\n", + " \n", + " ax[1].plot(nstep, planet2_chain1)\n", + " ax[1].plot(nstep, planet2_chain2)\n", + " \n", + " if plot_truth:\n", + " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", + " ax[1].set_title(\"Earth\", fontsize = 18)\n", + " ax[1].legend(fontsize = 18)\n", + "\n", + " fig.suptitle(param, fontsize = 18)\n", + " fig.tight_layout()\n", + " plt.show()\n", + "\n", + "\n", + "\n", + "\n", + " return \"plotting complete\"\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "def a_from_Kepler3(period, M_tot):\n", + " period = period*86400 #days to seconds\n", + " \n", + "\n", + " M_tot = M_tot*M_sun.value #solar masses to kg\n", + " \n", + " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", + " \n", + " a = a3**(1/3)\n", + " \n", + " \n", + " \n", + " a = a * 6.68459*10**(-12.) # meters to AU\n", + " \n", + " return(a) #in AUs\n", + "\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "id": "4ca33874", + "metadata": {}, + "source": [ + "## Check the convergence diagnostics:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cfd9e0cd", + "metadata": {}, + "outputs": [], + "source": [ + "import arviz as az\n", + "\n", + "az.summary(trace, var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"m_planet\",\n", + " \"Omega\",\n", + " \"omega\",\n", + " \"incl\"])\n" + ] + }, + { + "cell_type": "markdown", + "id": "6d8c580e", + "metadata": {}, + "source": [ + "#### Great, the r_hat statistitcs are all 1.0, suggesting that the chains are converged...\n", + "\n", + "#### ...we can also check this by eye by looking at the corner plots and the posterior chains!" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6775ae19", + "metadata": {}, + "outputs": [], + "source": [ + "import corner\n", + "\n", + "make_plots(trace, planet_params, 2)\n" + ] + }, + { + "cell_type": "markdown", + "id": "dd891787", + "metadata": {}, + "source": [ + "## Now, let's plot the MCMC posterior and simulated RV and Astrometry data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13742966", + "metadata": {}, + "outputs": [], + "source": [ + "import pandas as pd\n", + "#save simulated data as dataframe \n", + "simulated_data_dic = {\n", + " \"times_rv_observed\": x_rv,\n", + " \"rv_observed\": y_rv,\n", + " \"rv_err_observed\": y_rv_err,\n", + "\n", + " \"times_rv_orbit\": times_rv,\n", + " \"rv_orbit\": rv_orbit,\n", + " \"rv_orbit_sum\": rv_orbit_sum,\n", + "\n", + " \"times_astrometry_observed\": x_astrometry,\n", + " \"ra_observed\": ra_data,\n", + " \"dec_observed\": dec_data,\n", + " \"ra_err_observed\": ra_err,\n", + " \"dec_err_observed\": dec_err,\n", + "\n", + " \"times_astrometry_orbit\": times_astrometry,\n", + " \"ra_orbit\": ra_orbit,\n", + " \"ra_orbit_sum\": ra_orbit_sum,\n", + " \"dec_orbit\": dec_orbit,\n", + " \"dec_orbit_sum\": dec_orbit_sum\n", + " }\n", + "\n", + "\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", + "simulated_data = simulated_data.transpose()\n", + "simulated_data = simulated_data.fillna(value=np.nan)\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e4669aae", + "metadata": {}, + "outputs": [], + "source": [ + "def remove_nans(array):\n", + " array = array[~pd.isnull(array)]\n", + " \n", + " return array\n", + "\n", + "\n", + "\n", + "\n", + "def import_orbit_data(orbit_row):\n", + " import math\n", + " \n", + " orbit_row = remove_nans(orbit_row)\n", + " orbit_array = []\n", + " for row in orbit_row:\n", + " orbit_array.append([np.array(row[1]), row[0]])\n", + "\n", + " orbit_array = np.array(orbit_array)\n", + " \n", + " return orbit_array\n", + "\n", + "\n", + "\n", + "\n", + "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", + " import matplotlib.pyplot as plt\n", + "\n", + " \n", + " \n", + " #simulated RV observations\n", + " x_rv = simulated_data['times_rv_observed'].values\n", + " y_rv = simulated_data['rv_observed'].values\n", + " y_rv_err = simulated_data['rv_err_observed'].values\n", + " \n", + " #simulated RV full orbits\n", + " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", + " y_rv_orbit = simulated_data['rv_orbit'].values\n", + " \n", + " #simulated astrometric observations\n", + " x_astrometry = simulated_data['times_astrometry_observed'].values\n", + " y_ra = simulated_data['ra_observed'].values\n", + " y_dec = simulated_data['dec_observed'].values\n", + " y_ra_err = simulated_data['ra_err_observed'].values\n", + " y_dec_err = simulated_data['dec_err_observed'].values\n", + " \n", + " #simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", + " y_ra_orbit = simulated_data['ra_orbit'].values\n", + " y_dec_orbit = simulated_data['dec_orbit'].values\n", + " \n", + " \n", + " #remove any nans from simulated data (created because of mismatching lengths)\n", + " x_rv = remove_nans(x_rv)\n", + " y_rv = remove_nans(y_rv)\n", + " y_rv_err = remove_nans(y_rv_err)\n", + " x_rv_orbit = remove_nans(x_rv_orbit)\n", + " y_rv_orbit = import_orbit_data(y_rv_orbit)\n", + " x_astrometry = remove_nans(x_astrometry)\n", + " y_ra = remove_nans(y_ra)\n", + " y_dec = remove_nans(y_dec)\n", + " y_ra_err = remove_nans(y_ra_err)\n", + " y_dec_err = remove_nans(y_dec_err)\n", + " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", + " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", + " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", + " \n", + " \n", + " \n", + " \n", + " # make a fine grid that spans the observation window for plotting purposes\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", + "\n", + " # for predicted orbits\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " # add NaN for plotting purpose on true model\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", + " \n", + " \n", + " \n", + "\n", + "\n", + " \n", + " \n", + " figsAs = []\n", + " figsRV = []\n", + " for n in range(0, nplanets):\n", + " \n", + " ######\n", + " #make rv plots\n", + " \n", + " \n", + " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", + " ax0, ax1 = ax[0], ax[1]\n", + " \n", + " # Get the posterior median orbital parameters\n", + " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", + " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", + "\n", + " # Compute the median of posterior estimate the other planet. Then we can remove\n", + " # this from the data to plot just the planet we care about.\n", + " other = np.median(\n", + " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", + " )\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # Plot the data\n", + " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", + "\n", + " pred = np.percentile(\n", + " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", + " [16, 50, 84],\n", + " axis=(0, 1),\n", + " )\n", + " \n", + " med_rv = []\n", + " for an_x in x_rv:\n", + " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", + "\n", + "\n", + "\n", + " \n", + " # plot the MCMC model\n", + " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", + " art = ax0.fill_between(\n", + " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", + " art.set_edgecolor(\"none\")\n", + " \n", + " # plot the true RA model\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " ax0.legend(fontsize=9, loc=1)\n", + " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", + " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", + "\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " #plot residuals\n", + " ax1.axhline(0.0, color=\"k\")\n", + " ax1.errorbar(\n", + " x_rv, y_rv - other - med_rv,\n", + " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", + " )\n", + " \n", + " \n", + " figRV.tight_layout()\n", + " figRV.subplots_adjust(hspace=0.5)\n", + " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", + " figRV.tight_layout()\n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " figRV.show()\n", + " ################\n", + " #plot astrometry\n", + " \n", + " \n", + "\n", + " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", + " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", + "\n", + " # Get the posterior median orbital parameters\n", + " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", + " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", + "\n", + " # Compute the median of posterior estimate the other planet. Then we can remove\n", + " # this from the data to plot just the planet we care about.\n", + " other = np.median(\n", + " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", + " )\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # plot the DEC data\n", + " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", + " # plot the MCMC DEC model\n", + " pred = np.percentile(\n", + " trace.posterior[\"dec_fine\"].values[:, :, :, n],\n", + " [16, 50, 84],\n", + " axis=(0, 1),\n", + " )\n", + " \n", + " med_dec = []\n", + " for an_x in x_astrometry:\n", + " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", + " \n", + "\n", + "\n", + " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", + " art = ax0.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", + " art.set_edgecolor(\"none\")\n", + "\n", + " # plot the true DEC model\n", + " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + " \n", + " \n", + " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", + " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", + " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + "\n", + "\n", + "\n", + "\n", + " ax0.legend(fontsize=9, loc=1)\n", + " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", + " \n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", + " \n", + " \n", + " #plot residuals \n", + " ax1.axhline(0.0, color=\"k\")\n", + " ax1.errorbar(\n", + " x_astrometry, y_dec - other - med_dec,\n", + " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", + " )\n", + "\n", + "\n", + "\n", + " # Get the posterior median orbital parameters\n", + " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", + " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", + "\n", + " # Compute the median of posterior estimate the other planet. Then we can remove\n", + " # this from the data to plot just the planet we care about.\n", + " other = np.median(\n", + " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", + " )\n", + " \n", + "\n", + "\n", + "\n", + "\n", + "\n", + " # Plot the RA data\n", + " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", + "\n", + " pred = np.percentile(\n", + " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", + " [16, 50, 84],\n", + " axis=(0, 1), \n", + " )\n", + " \n", + " \n", + " med_ra = []\n", + " for an_x in x_astrometry:\n", + " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", + "\n", + "\n", + " # plot the MCMC RA model\n", + " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", + " art = ax2.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", + " art.set_edgecolor(\"none\")\n", + " \n", + " # plot the true RA model\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + " \n", + " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", + " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", + " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + "\n", + "\n", + " \n", + " ax2.legend(fontsize=9, loc=1)\n", + " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", + " \n", + " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", + "\n", + " \n", + " #plot residuals\n", + " ax3.axhline(0.0, color=\"k\")\n", + " ax3.errorbar(\n", + " x_astrometry, y_ra - other - med_ra,\n", + " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", + " )\n", + " \n", + " figAs.tight_layout()\n", + " figAs.subplots_adjust(wspace=.2)\n", + " \n", + " \n", + " figAs.show()\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", + " return None" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3941ebb", + "metadata": {}, + "outputs": [], + "source": [ + "make_rv_astrometry_plots(simulated_data, trace, 2)\n" + ] + }, + { + "cell_type": "markdown", + "id": "1f6c90cd", + "metadata": {}, + "source": [ + "### And lastly, check how well our MCMC median +/- 1 sigma results agree with the simulated values" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "672620cf", + "metadata": {}, + "outputs": [], + "source": [ + "def load_posterior_params(trace, nplanets):\n", + "\n", + " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", + "\n", + "\n", + " params_earth = defaultdict(list)\n", + " params_jup = defaultdict(list)\n", + " params_earth_err = defaultdict(list)\n", + " params_jup_err = defaultdict(list)\n", + "\n", + " table_params_all = []\n", + " posterior_meds_all = []\n", + " posterior_errs_all = []\n", + " for ii in range(0, nplanets):\n", + " table_params = defaultdict(list)\n", + " posterior_meds = defaultdict(list)\n", + " posterior_errs = defaultdict(list)\n", + " for param in parameters:\n", + "\n", + "\n", + " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", + "\n", + " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", + "\n", + "\n", + "\n", + " planet_err = [\n", + " planet_med - planet_quantile[0],\n", + " planet_quantile[1] - planet_med \n", + " ]\n", + "\n", + "\n", + "\n", + " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + "\n", + " posterior_meds[param] = np.round(planet_med, 3)\n", + " posterior_errs[param] = np.round(planet_err, 3)\n", + " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", + " str(np.round(planet_err[1], 3)) + \" -\" + \n", + " str(np.round(planet_err[0], 3)) + \")\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + " table_params_all.append(table_params)\n", + " posterior_meds_all.append(posterior_meds)\n", + " posterior_errs_all.append(posterior_errs)\n", + " #print(\"\")\n", + " #print(\"\")\n", + "\n", + "\n", + "\n", + " return table_params_all, posterior_meds_all, posterior_errs_all\n", + "\n", + "\n", + "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", + " posterior_meds, posterior_errs, nplanets):\n", + "\n", + " import matplotlib\n", + " import matplotlib.pyplot as plt\n", + " \n", + " \n", + " input_params = np.array(input_params)\n", + " input_params.T[-1] = input_params.T[-1]*332946.07\n", + "\n", + " comparisons = []\n", + " for ii in range(0, nplanets):\n", + " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", + " plt.subplots_adjust(hspace=1)\n", + " \n", + " #### add comparison plots\n", + " \n", + " \n", + " input_mass = input_params[ii][-1]\n", + " input_period = input_params[ii][0]\n", + " \n", + " xs = posterior_meds[ii]['m_planet']\n", + " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", + "\n", + " ys = posterior_meds[ii]['P']\n", + " ys_err = np.array([posterior_errs[ii]['P']]).T\n", + "\n", + "\n", + "\n", + " color = '#1c245f'\n", + "\n", + " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", + " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", + " linewidth = 3, alpha = 0.7)\n", + "\n", + " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", + "\n", + "\n", + " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", + " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", + "\n", + "\n", + " fig.suptitle('input and posterior parameters', fontsize = 27)\n", + "\n", + " axs[0].legend(fontsize = 13)\n", + " \n", + " \n", + " \n", + " #### add comparison tables\n", + "\n", + "\n", + " sim_params = input_params[ii]\n", + " post_params = posterior_orbit_params[ii]\n", + " data = np.column_stack((np.array(keys_orbit_params), \n", + " np.round(np.array(sim_params),3),\n", + " np.array(list(post_params.values()))))\n", + "\n", + " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", + "\n", + " df = pd.DataFrame(data, columns=labels)\n", + "\n", + " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", + " colColours =[\"#D7E5F0\"]*3, loc='center', \n", + " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", + " #for (row, col), cell in table.get_celld().items():\n", + " # if (row == 0):\n", + " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", + "\n", + "\n", + " table.auto_set_column_width(col=labels)\n", + "\n", + " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " axs[1].set_axis_off()\n", + "\n", + "\n", + " table.auto_set_font_size(False)\n", + " table.set_fontsize(13)\n", + "\n", + " \n", + " \n", + " \n", + "\n", + "\n", + "\n", + " comparisons.append(fig)\n", + " \n", + " \n", + " \n", + " \n", + " return comparisons" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "552e3f12", + "metadata": {}, + "outputs": [], + "source": [ + "from collections import defaultdict\n", + "\n", + "\n", + "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", + "load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", + "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", + " posterior_meds_all, posterior_errs_all, 2)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c491ea04", + "metadata": {}, + "outputs": [], + "source": [] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "96bfdecb", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.9.7" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From aa934327e8982ad8ab9f5cfb62662a58d5e35a74 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 9 Feb 2023 22:55:34 +0000 Subject: [PATCH 2/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/astrometric.ipynb | 1 - docs/tutorials/astrometry_rv_case_study.ipynb | 1740 +++++++++-------- docs/tutorials/eb.ipynb | 2 - docs/tutorials/joint.ipynb | 17 +- docs/tutorials/lc-gp-transit.ipynb | 4 - docs/tutorials/lc-multi.ipynb | 2 - docs/tutorials/quick-tess.ipynb | 1 - docs/tutorials/rv-multi.ipynb | 1 - docs/tutorials/rv.ipynb | 1 - docs/tutorials/stellar-variability.ipynb | 1 - docs/tutorials/tess.ipynb | 9 +- docs/tutorials/transit.ipynb | 1 - docs/tutorials/ttv.ipynb | 1 - 13 files changed, 925 insertions(+), 856 deletions(-) diff --git a/docs/tutorials/astrometric.ipynb b/docs/tutorials/astrometric.ipynb index 0e32ab93..d871c170 100644 --- a/docs/tutorials/astrometric.ipynb +++ b/docs/tutorials/astrometric.ipynb @@ -286,7 +286,6 @@ "\n", "def get_model(parallax=None):\n", " with pm.Model() as model:\n", - "\n", " if parallax is None:\n", " # Without an actual parallax measurement, we can model the orbit in units of arcseconds\n", " # by providing a fake_parallax and conversion constant\n", diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb index 4e43e340..7dcc7410 100644 --- a/docs/tutorials/astrometry_rv_case_study.ipynb +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -12,7 +12,9 @@ "cell_type": "code", "execution_count": null, "id": "4abd3961", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -20,8 +22,7 @@ "\n", "\n", "xo.utils.docs_setup()\n", - "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", - "\n" + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")" ] }, { @@ -29,7 +30,7 @@ "id": "9a87ff47", "metadata": {}, "source": [ - "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023.\n", "\n", "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", "\n", @@ -49,12 +50,14 @@ { "cell_type": "markdown", "id": "93ceb2a9", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## (Simulated) Data\n", "\n", "\n", - "First, we need to simulate our dataset. \n", + "First, we need to simulate our dataset.\n", "\n", "\n", "### Let's define 2 helpful functions in order to simulate our data...." @@ -68,27 +71,24 @@ "outputs": [], "source": [ "def create_orbit(n_planets, orbit_params):\n", - " \n", - " '''\n", + " \"\"\"\n", " inputs:\n", " - n_planets = number of planets\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", - " '''\n", - " \n", + " - m_planet = mass of planet in Earth masses\n", + " \"\"\"\n", "\n", - " \n", " period = []\n", " ecc = []\n", - " Tper = [] \n", - " omega = [] \n", - " Omega = [] \n", + " Tper = []\n", + " omega = []\n", + " Omega = []\n", " incl = []\n", " m_planet = []\n", " for ii in range(0, n_planets):\n", @@ -101,15 +101,15 @@ " m_planet.append(orbit_params[ii][6])\n", "\n", " orbit = xo.orbits.KeplerianOrbit(\n", - " period = period,\n", - " ecc = ecc, \n", - " t_periastron = Tper, \n", - " omega = omega, \n", - " Omega = Omega, \n", - " incl = incl,\n", - " m_planet = m_planet) \n", - " \n", - " \n", + " period=period,\n", + " ecc=ecc,\n", + " t_periastron=Tper,\n", + " omega=omega,\n", + " Omega=Omega,\n", + " incl=incl,\n", + " m_planet=m_planet,\n", + " )\n", + "\n", " return orbit" ] }, @@ -120,24 +120,24 @@ "metadata": {}, "outputs": [], "source": [ - "np.random.seed(979) \n", + "np.random.seed(979)\n", + "\n", "\n", "def simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra,\n", " sigma_dec,\n", " plx,\n", " orbit_params,\n", - " times_observed_rv = None,\n", - " times_observed_astrometry = None,\n", - " t_dur_rv = None, \n", - " n_obs_rv = None,\n", - " t_dur_astrometry = None, \n", - " n_obs_astrometry = None \n", - " ):\n", - "\n", - " '''\n", + " times_observed_rv=None,\n", + " times_observed_astrometry=None,\n", + " t_dur_rv=None,\n", + " n_obs_rv=None,\n", + " t_dur_astrometry=None,\n", + " n_obs_astrometry=None,\n", + "):\n", + " \"\"\"\n", " inputs:\n", " - orbit = keplerian orbit generated by exoplanet code\n", " - n_planets = number of planets\n", @@ -147,230 +147,231 @@ " - plx = parallax of system\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", + " - m_planet = mass of planet in Earth masses\n", " - times_observed_rv = (default=None) the observed times for RV observations\n", " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", - " \n", - " \n", - " '''\n", - " \n", + "\n", + "\n", + " \"\"\"\n", + "\n", " if n_planets == 1:\n", " orbit_params = [orbit_params]\n", - " \n", - " \n", - " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", - " if times_observed_rv != None:\n", - " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", - "\n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined\n", + "\n", + " # if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(\n", + " np.min(times_observed_rv), np.max(times_observed_rv), 10000\n", + " )\n", + "\n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_rv = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_rv:\n", " start_time_rv = orbit_params[ii][2]\n", "\n", + " # define RV observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv + t_dur_rv, 10000)\n", + " times_observed_rv = np.linspace(\n", + " start_time_rv, start_time_rv + t_dur_rv, n_obs_rv\n", + " )\n", "\n", - " \n", - " #define RV observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", - " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", - " \n", - "\n", + " # if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(\n", + " np.min(times_observed_astrometry),\n", + " np.max(times_observed_astrometry),\n", + " 10000,\n", + " )\n", "\n", - " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", - " if times_observed_astrometry != None:\n", - " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", - " \n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined \n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_astrometry = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_astrometry:\n", - " start_time_astrometry = orbit_params[ii][2] \n", - "\n", - "\n", - " #define astrometry observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", - " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", - "\n", + " start_time_astrometry = orbit_params[ii][2]\n", + "\n", + " # define astrometry observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " 10000,\n", + " )\n", + " times_observed_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " n_obs_astrometry,\n", + " )\n", "\n", - " \n", " orbit = create_orbit(n_planets, orbit_params)\n", - " \n", + "\n", " rv = orbit.get_radial_velocity(times_rv)\n", " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", - " \n", + "\n", " rv_orbit = rv.eval()\n", "\n", " if n_planets > 1:\n", - " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", + " rv_orbit_sum = np.sum(rv_orbit, axis=1)\n", " else:\n", " rv_orbit_sum = rv_orbit\n", - " \n", "\n", " rv_observed = rv_observed.eval()\n", " if n_planets > 1:\n", - " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", + " rv_observed_sum = np.sum(rv_observed, axis=1)\n", " else:\n", " rv_observed_sum = rv_observed\n", "\n", - "\n", - " #-----------\n", - " #-----------\n", + " # -----------\n", + " # -----------\n", "\n", " # determine and print the star position at desired times\n", " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", "\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", - " #pos = tt.sum(pos, axis=-1)\n", - "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " rho = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", + "\n", " rho, theta = rho.eval(), theta.eval()\n", - " \n", - " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", + " # rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", " rho_orbit = rho\n", - " \n", + "\n", " if n_planets > 1:\n", - " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", + " rho_orbit_sum = np.sum(rho_orbit, axis=1)\n", " else:\n", " rho_orbit_sum = rho_orbit\n", "\n", " theta_orbit = theta\n", - " \n", + "\n", " if n_planets > 1:\n", - " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", + " theta_orbit_sum = np.sum(theta_orbit, axis=1)\n", " else:\n", " theta_orbit_sum = theta_orbit\n", - " \n", - " \n", + "\n", " # when summing over theta, position angle, we have to careful because position\n", - " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract\n", " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", " # in the correct range. Be careful though if modeling more than 2 planets, this\n", " # doesn't completely solve the problem!\n", - " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", - " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", - "\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2 * np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2 * np.pi\n", "\n", " # determine and print the star position at desired times\n", - " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", - "\n", - "\n", - " #pos = tt.sum(pos, axis=-1)\n", + " pos_observed = theano.function(\n", + " [], orbit.get_star_position(times_observed_astrometry, plx)\n", + " )()\n", "\n", - " x_obs,y_obs,z_obs = pos_observed\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", + " x_obs, y_obs, z_obs = pos_observed\n", "\n", " # calculate rho and theta\n", - " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", - " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", - " \n", - " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs**2 + y_obs**2)) # arcsec\n", + " theta_observed = tt.squeeze(\n", + " tt.arctan2(y_obs, x_obs)\n", + " ) # radians between [-pi, pi]\n", "\n", + " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", "\n", - " \n", " if n_planets > 1:\n", - " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", + " rho_observed_sum = np.sum(rho_observed, axis=1)\n", " else:\n", " rho_observed_sum = rho_observed\n", "\n", - " \n", " if n_planets > 1:\n", - " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", + " theta_observed_sum = np.sum(theta_observed, axis=1)\n", " else:\n", " theta_observed_sum = theta_observed\n", - " \n", - " \n", + "\n", " #######\n", " #######\n", " #######\n", "\n", " # convert rho and theta into ra and dec\n", - " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", "\n", " if n_planets > 1:\n", - " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", + " ra_orbit_sum = np.sum(ra_orbit, axis=1)\n", " else:\n", " ra_orbit_sum = ra_orbit\n", "\n", - " \n", " if n_planets > 1:\n", - " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", + " dec_orbit_sum = np.sum(dec_orbit, axis=1)\n", " else:\n", " dec_orbit_sum = dec_orbit\n", "\n", - " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", "\n", - "\n", - "\n", - "\n", - " \n", - " #----------- \n", - " #----------- \n", - " #----------- simulate data\n", - " #----------- \n", - " #----------- \n", - "\n", + " # -----------\n", + " # -----------\n", + " # ----------- simulate data\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", " rv_sim = rv_observed\n", - " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", + " rv_sim_sum = np.sum(rv_observed, axis=1) + np.random.normal(\n", + " 0, sigma_rv, len(rv_observed)\n", + " )\n", " else:\n", " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", " rv_sim_sum = rv_sim\n", "\n", - " #-----------\n", - " #-----------\n", - "\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", - " ra_sim = ra_observed \n", - " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", - " \n", + " ra_sim = ra_observed\n", + " ra_sim_sum = np.sum(ra_sim, axis=1) + np.random.normal(\n", + " 0, sigma_ra, len(ra_observed)\n", + " )\n", + "\n", " else:\n", " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", " ra_sim_sum = ra_sim\n", - " \n", - " \n", - " \n", + "\n", " if n_planets > 1:\n", - " dec_sim = dec_observed \n", - " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", - " \n", + " dec_sim = dec_observed\n", + " dec_sim_sum = np.sum(dec_sim, axis=1) + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", + "\n", " else:\n", - " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " dec_sim = dec_observed + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", " dec_sim_sum = dec_sim\n", "\n", - " \n", - " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", + " times = [\n", + " times_rv,\n", + " times_observed_rv,\n", + " times_astrometry,\n", + " times_observed_astrometry,\n", + " ]\n", " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", - " \n", - " \n", + "\n", " return [times, rv_results, ra_results, dec_results]" ] }, @@ -398,10 +399,20 @@ "source": [ "import numpy as np\n", "\n", - "parallax = 0.1 #as \n", - "n_planets=2\n", - "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", - " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" + "parallax = 0.1 # as\n", + "n_planets = 2\n", + "planet_params = [\n", + " [365.0, 0.0167, 100.0, 1.7959438003021653, 0.0, np.pi / 4, 3e-06],\n", + " [\n", + " 4327.0,\n", + " 0.0484,\n", + " 500.0,\n", + " -1.4957471689591395,\n", + " 1.7523105690023069,\n", + " np.pi / 4 + 0.022759,\n", + " 0.000955,\n", + " ],\n", + "]" ] }, { @@ -422,28 +433,28 @@ "cell_type": "code", "execution_count": null, "id": "55938572", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "from aesara_theano_fallback import aesara as theano\n", "import theano.tensor as tt\n", "\n", "\n", - "#add gaia observing times\n", + "# add gaia observing times\n", "times_observed_astrometry_gaia = []\n", "t_0 = 0\n", - "for ii in range(t_0, t_0+3600):\n", - " if ii % (3600/100) == 0:\n", + "for ii in range(t_0, t_0 + 3600):\n", + " if ii % (3600 / 100) == 0:\n", " times_observed_astrometry_gaia.append(ii)\n", "\n", "\n", - "\n", - "#add THE observing times\n", + "# add THE observing times\n", "times_observed_rv = []\n", - "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "t_0 = 3300 # Gaia first light (2014) ~9 years before THE first expected light (2023)\n", "add_data = True\n", - "for ii in range(t_0, t_0+3600):\n", - "\n", + "for ii in range(t_0, t_0 + 3600):\n", " if ii % 180 == 0:\n", " if add_data:\n", " add_data = False\n", @@ -454,72 +465,66 @@ " times_observed_rv.append(ii)\n", "\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "sigma_rv = 0.3\n", "\n", - "sigma_ra_gaia = 6e-5 #60 micro-as\n", - "sigma_dec_gaia = 6e-5 #60 micro-as\n", - "\n", + "sigma_ra_gaia = 6e-5 # 60 micro-as\n", + "sigma_dec_gaia = 6e-5 # 60 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_gaia,\n", " sigma_dec_gaia,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_gaia\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_gaia,\n", + ")\n", "\n", "\n", - "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", - "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", - "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", - "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", + "[\n", + " [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + " [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + " [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + " [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum],\n", + "] = (times, rv_results, ra_results, dec_results)\n", "\n", "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", "\n", "\n", - "\n", - "\n", - "\n", "# we assume that the first roman observation\n", "# will occur ~5 years after the last Gaia observation\n", - "t_1 = times_observed_astrometry_gaia[-1]+1800\n", + "t_1 = times_observed_astrometry_gaia[-1] + 1800\n", "\n", "# add roman observing times\n", "times_observed_astrometry_roman = []\n", - "for ii in range(t_1, t_1+(10*365)):\n", + "for ii in range(t_1, t_1 + (10 * 365)):\n", " if ii % 90 == 0:\n", - " times_observed_astrometry_roman.append(ii) \n", - "\n", + " times_observed_astrometry_roman.append(ii)\n", "\n", - "sigma_ra_roman = 5e-6 #5 micro-as\n", - "sigma_dec_roman = 5e-6 #5 micro-as\n", "\n", + "sigma_ra_roman = 5e-6 # 5 micro-as\n", + "sigma_dec_roman = 5e-6 # 5 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_roman,\n", " sigma_dec_roman,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_roman\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_roman,\n", + ")\n", "\n", "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", "\n", - "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", + "times_observed_astrometry = np.append(\n", + " times_observed_astrometry, times[3], axis=0\n", + ")\n", "\n", "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", @@ -532,7 +537,7 @@ "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", "\n", "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", - "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)" ] }, { @@ -552,36 +557,40 @@ "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib\n", - "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", - "matplotlib.rc('text', usetex=True)\n", + "\n", + "matplotlib.rc(\"font\", **{\"family\": \"serif\", \"serif\": [\"Computer Modern\"]})\n", + "matplotlib.rc(\"text\", usetex=True)\n", + "\n", "\n", "def plot_rv_signal(\n", - " n_planets, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv):\n", - " \n", - " \n", - " \n", - " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", - " fig.suptitle(\"RV Signal\", fontsize = 45)\n", + " n_planets,\n", + " rv_orbit,\n", + " rv_orbit_sum,\n", + " rv_sim,\n", + " rv_sim_sum,\n", + " times_rv,\n", + " times_observed_rv,\n", + "):\n", + " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", + " fig.suptitle(\"RV Signal\", fontsize=45)\n", " ax1 = ax\n", - " \n", - " \n", - " \n", "\n", - " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", + " ax1.plot(\n", + " times_observed_rv,\n", + " rv_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " alpha=0.3,\n", + " )\n", "\n", - " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", - " ax1.legend(fontsize = 27, loc = 2)\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize=27)\n", + " ax1.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -593,53 +602,55 @@ "outputs": [], "source": [ "def plot_astrometry_signal(\n", - " n_planets, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " n_planets,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry):\n", - " \n", - " \n", - "\n", - " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", - " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", + " times_observed_astrometry,\n", + "):\n", + " fig, ax = plt.subplots(2, 1, figsize=[13, 10], sharex=\"col\", sharey=\"row\")\n", + " fig.suptitle(\"Astrometric Signal\", fontsize=45)\n", " ax2 = ax[0]\n", " ax3 = ax[1]\n", - " \n", - "\n", "\n", - " \n", - " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.plot(\n", + " times_observed_astrometry,\n", + " ra_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", " for tick in ax3.get_xticklabels():\n", " tick.set_rotation(30)\n", - " ax2.legend(fontsize = 27, loc = 2)\n", - " \n", - " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", - " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax3.legend(fontsize = 27, loc = 2)\n", - "\n", - "\n", + " ax2.legend(fontsize=27, loc=2)\n", + "\n", + " ax3.plot(\n", + " times_observed_astrometry,\n", + " dec_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax3.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", "\n", - "\n", - "\n", " #######\n", " #######\n", " #######\n", "\n", - "\n", - " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", + " fig, ax = plt.subplots(1, 1, figsize=[10, 10])\n", "\n", " # plot simulated data\n", - " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", + " ax.plot(ra_sim_sum, dec_sim_sum, \"o\", color=\"k\", label=\"combined signal\")\n", "\n", " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", @@ -649,7 +660,7 @@ " ax.set_title(\"initial orbit\")\n", " ax.legend()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -665,46 +676,46 @@ "cell_type": "code", "execution_count": null, "id": "82e18e64", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_rv_signal(\n", - " 2, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv)\n", - " \n", - " \n" + " 2, rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum, times_rv, times_observed_rv\n", + ")" ] }, { "cell_type": "code", "execution_count": null, "id": "623b9c2d", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_astrometry_signal(\n", - " 2, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " 2,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry)" + " times_observed_astrometry,\n", + ")" ] }, { "cell_type": "markdown", "id": "99110cc0", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Ok, now we have a simulated dataset -- let's model it!\n", "\n", @@ -725,9 +736,7 @@ " print(\"------------\")\n", "\n", " with pm.Model() as model:\n", - " \n", - " \n", - " #log normal prior on period around estimates\n", + " # log normal prior on period around estimates\n", " logP = pm.Uniform(\n", " \"logP\",\n", " lower=0,\n", @@ -735,37 +744,33 @@ " shape=2,\n", " testval=np.log(periods),\n", " )\n", - " \n", - "\n", "\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", "\n", - "\n", " ## wide uniform prior on t_periastron\n", " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", - " \n", - " \n", + "\n", " # Wide normal prior for semi-amplitude\n", - " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", - " \n", + " logK = pm.Uniform(\n", + " \"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks)\n", + " )\n", + "\n", " K = pm.Deterministic(\"K\", tt.exp(logK))\n", - " \n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", "\n", - "\n", - "\n", - " \n", " # Jitter & a quadratic RV trend\n", - " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + " # logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", "\n", " # Then we define the orbit\n", - " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", + " orbit = xo.orbits.KeplerianOrbit(\n", + " period=P, t_periastron=tperi, ecc=ecc, omega=omega\n", + " )\n", "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", @@ -776,8 +781,6 @@ " # Sum over planets and add the background to get the full model\n", " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", - " \n", - "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", "\n", @@ -786,19 +789,17 @@ "\n", " # Finally add in the observation model. This next line adds a new contribution\n", " # to the log probability of the PyMC3 model\n", - " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " # err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", " err = y_rv_err\n", " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", "\n", - "\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", " map_soln = pmx.optimize(start=map_soln)\n", "\n", - "\n", - " #return the max a-posteriori solution\n", + " # return the max a-posteriori solution\n", " return map_soln" ] }, @@ -806,16 +807,18 @@ "cell_type": "code", "execution_count": null, "id": "4d6faa40", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import pymc3 as pm\n", "import pymc3_ext as pmx\n", "\n", - "np.random.seed(979) \n", + "np.random.seed(979)\n", "\n", "\n", - "#rename variables in more consistent way for modeling\n", + "# rename variables in more consistent way for modeling\n", "x_rv = np.array(times_observed_rv)\n", "y_rv = rv_sim_sum\n", "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", @@ -829,25 +832,29 @@ "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", "\n", "\n", - " \n", - " \n", "# make a fine grid that spans the observation window for plotting purposes\n", - "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + "t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + ")\n", "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", "# for predicted orbits\n", - "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + ")\n", "\n", "\n", - "#minimize on RV data\n", + "# minimize on RV data\n", "periods_real = [4327.0, 365.0]\n", "periods_guess = []\n", "\n", - "#let's assume that THE can recover the period within 10% of the true value\n", + "# let's assume that THE can recover the period within 10% of the true value\n", "for per in periods_real:\n", - " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", + " periods_guess.append(\n", + " per + np.random.uniform(low=-0.1 * per, high=0.1 * per)\n", + " )\n", "\n", - "print('periods used as starting point = ', periods_guess)\n", + "print(\"periods used as starting point = \", periods_guess)\n", "periods_guess = np.array(periods_guess)\n", "\n", "\n", @@ -859,7 +866,9 @@ { "cell_type": "markdown", "id": "efc4dcc2", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Now let's define out joint astrometry and radial velocity model:" ] @@ -871,25 +880,40 @@ "metadata": {}, "outputs": [], "source": [ - "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", - " m_sun = 333030 #earth masses\n", - " \n", - " P_RV = np.array(rv_map_soln['P'])\n", - " K_RV = np.array(rv_map_soln['K'])\n", - " tperi_RV = np.array(rv_map_soln['tperi'])\n", - " ecc_RV = np.array(rv_map_soln['ecc'])\n", - " omega_RV = np.array(rv_map_soln['omega'])\n", - " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", + "def minimize_both(\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", + "):\n", + " m_sun = 333030 # earth masses\n", + "\n", + " P_RV = np.array(rv_map_soln[\"P\"])\n", + " K_RV = np.array(rv_map_soln[\"K\"])\n", + " tperi_RV = np.array(rv_map_soln[\"tperi\"])\n", + " ecc_RV = np.array(rv_map_soln[\"ecc\"])\n", + " omega_RV = np.array(rv_map_soln[\"omega\"])\n", + " min_masses_RV = (\n", + " xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value * 317.83\n", + " ) # in m_earth\n", " phase_RV = determine_phase(P_RV, tperi_RV)\n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - "\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", " print(\"RV Solutions\")\n", " print(\"------------\")\n", @@ -899,98 +923,106 @@ " print(\"eccentricity: \", ecc_RV)\n", " print(\"omega: \", omega_RV)\n", "\n", - " print('')\n", + " print(\"\")\n", " print(\"minimizing joint model solutions pre-MCMC\")\n", " print(\"------------\")\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", - " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", + " inc_test_vals = np.array(np.radians([5.0, 25.0, 45.0, 65.0, 85.0]))\n", " model, map_soln, logp = [], [], []\n", " for inc in inc_test_vals:\n", - " mass_test_vals = min_masses_RV/np.sin(inc)\n", - " print('')\n", + " mass_test_vals = min_masses_RV / np.sin(inc)\n", + " print(\"\")\n", " print(\"trying inclination = \" + str(np.degrees(inc)))\n", " print(\"mass test val = \" + str(mass_test_vals))\n", " print(\"------------\")\n", - " \n", - " \n", - "\n", - "\n", "\n", " def get_model():\n", " with pm.Model() as model:\n", - "\n", - "\n", " # Below we will run a version of this model where a measurement of parallax is provided\n", " # The measurement is in milliarcsec\n", " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", - " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", + " \"m_plx\", mu=parallax * 1000, sd=10, testval=parallax * 1000\n", " )\n", " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", "\n", - "\n", " # We expect the period to be around that found from just the RVs,\n", " # so we'll set a broad prior on logP\n", - " \n", + "\n", " logP = pm.Uniform(\n", - " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", + " \"logP\",\n", + " lower=0,\n", + " upper=np.log(2 * P_RV),\n", + " testval=np.log(P_RV),\n", + " shape=2,\n", " )\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", - " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", - " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", - " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecs = pmx.UnitDisk(\n", + " \"ecs\",\n", + " shape=(2, 2),\n", + " testval=np.array(\n", + " [\n", + " np.sqrt(ecc_RV) * np.cos(omega_RV),\n", + " np.sqrt(ecc_RV) * np.sin(omega_RV),\n", + " ]\n", + " ),\n", + " )\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", - " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", - "\n", - " \n", - " \n", + " xo.eccentricity.kipping13(\n", + " \"ecc_prior\", fixed=True, observed=ecc\n", + " )\n", "\n", " # Omegas are co-dependent, so sample them with variables Omega_plus\n", - " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and\n", " # Omega_minus is (Omega_0 - Omega_1)/2\n", - " \n", + "\n", " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", - " \n", - " \n", - " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", - " (Omega_plus - Omega_minus)] )\n", - " \n", "\n", - " Omega = pm.Deterministic(\"Omega\", Omega) \n", - " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", - " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", - " \n", + " Omega = tt.concatenate(\n", + " [(Omega_plus + Omega_minus), (Omega_plus - Omega_minus)]\n", + " )\n", "\n", + " Omega = pm.Deterministic(\"Omega\", Omega)\n", + " Omega_sum = pm.Deterministic(\n", + " \"Omega_sum\", ((Omega_plus) * 2) % np.pi\n", + " )\n", + " Omega_diff = pm.Deterministic(\n", + " \"Omega_diff\", ((Omega_minus) * 2) % np.pi\n", + " )\n", "\n", - " \n", " # For these orbits, it can also be better to fit for a phase angle\n", " # (relative to a reference time) instead of the time of periasteron\n", " # passage directly\n", " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", - " \n", "\n", - " \n", " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", - " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", + " log_m = pm.Uniform(\n", + " \"log_m\",\n", + " lower=-1,\n", + " upper=np.log(100 * mass_test_vals),\n", + " testval=np.log(mass_test_vals),\n", + " shape=2,\n", + " )\n", " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", - " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", - "\n", + " m_planet_fit = pm.Deterministic(\n", + " \"m_planet_fit\", m_planet / m_sun\n", + " )\n", "\n", - " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", + " cos_incl = pm.Uniform(\n", + " \"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2\n", + " )\n", " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", - " \n", - "\n", "\n", - " \n", - " \n", " # Set up the orbit\n", " orbit = xo.orbits.KeplerianOrbit(\n", " t_periastron=tperi,\n", @@ -999,13 +1031,10 @@ " ecc=ecc,\n", " omega=omega,\n", " Omega=Omega,\n", - " m_planet = m_planet_fit,\n", - " plx=plx\n", + " m_planet=m_planet_fit,\n", + " plx=plx,\n", " )\n", "\n", - "\n", - " \n", - " \n", " # Add a function for computing the full astrometry model\n", " def get_astrometry_model(t, name=\"\"):\n", " # First the astrometry induced by the planets\n", @@ -1013,54 +1042,56 @@ " # determine and print the star position at desired times\n", " pos = orbit.get_star_position(t, plx)\n", "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", - " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", - " \n", - " #rhos, thetas = get_star_relative_angles(t, plx)\n", - " \n", - " \n", - " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", - " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", - " \n", + " rhos = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", + " thetas = tt.squeeze(\n", + " tt.arctan2(y, x)\n", + " ) # radians between [-pi, pi]\n", + "\n", + " # rhos, thetas = get_star_relative_angles(t, plx)\n", + "\n", + " dec = pm.Deterministic(\n", + " \"dec\" + name, rhos * np.cos(thetas)\n", + " ) # X is north\n", + " ra = pm.Deterministic(\n", + " \"ra\" + name, rhos * np.sin(thetas)\n", + " ) # Y is east\n", + "\n", " # Sum over planets to get the full model\n", - " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", - " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", - " \n", + " dec_model = pm.Deterministic(\n", + " \"dec_model\" + name, tt.sum(dec, axis=-1)\n", + " )\n", + " ra_model = pm.Deterministic(\n", + " \"ra_model\" + name, tt.sum(ra, axis=-1)\n", + " )\n", "\n", - " \n", " return dec_model, ra_model\n", "\n", - " \n", " # Define the astrometry model at the observed times\n", " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", "\n", " # Also define the model on a fine grid as computed above (for plotting)\n", - " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", - "\n", - " \n", - "\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(\n", + " t_fine, name=\"_fine\"\n", + " )\n", "\n", - " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", - " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " # dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " # ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", " dec_tot_err = dec_err\n", " ra_tot_err = ra_err\n", "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec\n", + " pm.Normal(\n", + " \"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data\n", + " )\n", + " pm.Normal(\n", + " \"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data\n", + " )\n", "\n", - " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", - " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", - " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", - "\n", - "\n", - " \n", - " \n", " # ADD RV MODEL\n", "\n", - "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", " # First the RVs induced by the planets\n", @@ -1068,7 +1099,9 @@ " pm.Deterministic(\"vrad\" + name, vrad)\n", "\n", " # Sum over planets to get the full model\n", - " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", + " return pm.Deterministic(\n", + " \"rv_model\" + name, tt.sum(vrad, axis=-1)\n", + " )\n", "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", @@ -1084,24 +1117,20 @@ " # Optimize to find the initial parameters\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", - " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + " # map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", "\n", " map_soln = pmx.optimize(map_soln)\n", "\n", - "\n", " return model, map_soln\n", "\n", - "\n", - "\n", " a_model, a_map_soln = get_model()\n", - " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", - " print('log likelihood = ' + str(a_logp))\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis=0)\n", + " print(\"log likelihood = \" + str(a_logp))\n", "\n", " model.append(a_model)\n", " map_soln.append(a_map_soln)\n", " logp.append(a_logp)\n", "\n", - " \n", " best_index = 0\n", " for index in range(0, len(model)):\n", " if logp[index] >= logp[best_index]:\n", @@ -1111,31 +1140,26 @@ " the_map_soln = map_soln[best_index]\n", " the_logp = logp[best_index]\n", "\n", - "\n", - "\n", " return the_model, the_map_soln, the_logp\n", "\n", "\n", - "\n", "def determine_phase(P, t_periastron):\n", " phase = (2 * np.pi * t_periastron) / P\n", " return phase\n", "\n", "\n", - "\n", "def model_both(model, map_soln, tune_steps, draw_steps):\n", - " print('Joint RV + Astometry Minimization Solutions:')\n", + " print(\"Joint RV + Astometry Minimization Solutions:\")\n", " print(\"------------\")\n", "\n", - " print('m_planet: ' + str(map_soln['m_planet']))\n", - " print('P: ' + str(map_soln['P']))\n", - " print('incl: ' + str(map_soln['incl']))\n", - " print('Omega: ' + str(map_soln['Omega']))\n", - " print('tperi: ' + str(map_soln['tperi']))\n", - " print('ecc: ' + str(map_soln['ecc']))\n", - " print('omega: ' + str(map_soln['omega']))\n", - " print('plx: ' + str(map_soln['plx']))\n", - "\n", + " print(\"m_planet: \" + str(map_soln[\"m_planet\"]))\n", + " print(\"P: \" + str(map_soln[\"P\"]))\n", + " print(\"incl: \" + str(map_soln[\"incl\"]))\n", + " print(\"Omega: \" + str(map_soln[\"Omega\"]))\n", + " print(\"tperi: \" + str(map_soln[\"tperi\"]))\n", + " print(\"ecc: \" + str(map_soln[\"ecc\"]))\n", + " print(\"omega: \" + str(map_soln[\"omega\"]))\n", + " print(\"plx: \" + str(map_soln[\"plx\"]))\n", "\n", " with model:\n", " trace = pmx.sample(\n", @@ -1163,18 +1187,24 @@ "\n", "################\n", "################\n", - "#minimize on joint model\n", + "# minimize on joint model\n", "joint_model, joint_map_soln, joint_logp = minimize_both(\n", - " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", - " ra_data, ra_err, dec_data, dec_err, parallax\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", ")\n", "\n", "\n", - "\n", - "\n", "################\n", "################\n", - "#run full MCMC\n", + "# run full MCMC\n", "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" ] }, @@ -1193,143 +1223,168 @@ "metadata": {}, "outputs": [], "source": [ - "matplotlib.rc('text', usetex=False)\n", + "matplotlib.rc(\"text\", usetex=False)\n", + "\n", "\n", "def make_plots(trace, orbit_params, n_planets):\n", - " m_sun = 333030 #earth masses\n", + " m_sun = 333030 # earth masses\n", "\n", " [orbit_params_earth, orbit_params_jup] = orbit_params\n", "\n", - "\n", - " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", - " inclination_earth, m_earth] = orbit_params_earth\n", - "\n", - "\n", - " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", - " inclination_jup, m_jup] = orbit_params_jup\n", - "\n", - "\n", - " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", - " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", - " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", - " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", - " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", - "\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", - " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", - " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " [\n", + " P_earth,\n", + " e_earth,\n", + " Tper_earth,\n", + " omega_earth,\n", + " Omega_earth,\n", + " inclination_earth,\n", + " m_earth,\n", + " ] = orbit_params_earth\n", + "\n", + " [\n", + " P_jup,\n", + " e_jup,\n", + " Tper_jup,\n", + " omega_jup,\n", + " Omega_jup,\n", + " inclination_jup,\n", + " m_jup,\n", + " ] = orbit_params_jup\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0 + m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0 + m_jup)\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " Omega_earth - Omega_jup,\n", + " Omega_earth + Omega_jup,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " 0.1,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " ]\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " ]\n", "\n", " # plot the table summarizing MCMC results\n", " az.summary(\n", " trace,\n", - " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", - " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", + " var_names=[\n", + " \"P\",\n", + " \"tperi\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_plus\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"ecc\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " ],\n", " )\n", "\n", - "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", " # plot the corner plots\n", " _ = corner.corner(\n", - " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", - " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", - " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", - " truths = truths, truth_color = \"#03003a\"\n", + " trace,\n", + " var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " \"tperi\",\n", + " ],\n", + " quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True,\n", + " title_kwargs={\"fontsize\": 13},\n", + " truths=truths,\n", + " truth_color=\"#03003a\",\n", " )\n", "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the chains\n", " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", " for ii in range(0, len(parameters)):\n", " plot_truth = False\n", " param = parameters[ii]\n", - " \n", - " true_vals_earth = truth_chain_plot[2*ii]\n", - " true_vals_jup = truth_chain_plot[2*ii+1]\n", + "\n", + " true_vals_earth = truth_chain_plot[2 * ii]\n", + " true_vals_jup = truth_chain_plot[2 * ii + 1]\n", " plot_truth = True\n", - " \n", - " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", + "\n", + " fig, ax = plt.subplots(1, 2, figsize=(15, 3))\n", " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", - " \n", + "\n", " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", - " \n", - " \n", - " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", - " \n", - " \n", + "\n", + " nstep = np.arange(1, len(planet1_chain1) + 1, 1)\n", + "\n", " ax[0].plot(nstep, planet1_chain1)\n", " ax[0].plot(nstep, planet1_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", - " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", - " ax[0].legend(fontsize = 18)\n", - " \n", + " ax[0].axhline(y=true_vals_earth, color=\"r\", label=\"truth\")\n", + " ax[0].set_title(\"Jupiter\", fontsize=18)\n", + " ax[0].legend(fontsize=18)\n", + "\n", " ax[1].plot(nstep, planet2_chain1)\n", " ax[1].plot(nstep, planet2_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", - " ax[1].set_title(\"Earth\", fontsize = 18)\n", - " ax[1].legend(fontsize = 18)\n", + " ax[1].axhline(y=true_vals_jup, color=\"r\", label=\"truth\")\n", + " ax[1].set_title(\"Earth\", fontsize=18)\n", + " ax[1].legend(fontsize=18)\n", "\n", - " fig.suptitle(param, fontsize = 18)\n", + " fig.suptitle(param, fontsize=18)\n", " fig.tight_layout()\n", " plt.show()\n", "\n", + " return \"plotting complete\"\n", "\n", "\n", + "def a_from_Kepler3(period, M_tot):\n", + " period = period * 86400 # days to seconds\n", "\n", - " return \"plotting complete\"\n", - "\n", + " M_tot = M_tot * M_sun.value # solar masses to kg\n", "\n", + " a3 = (((constants.G.value) * M_tot) / (4 * np.pi**2)) * (period) ** 2.0\n", "\n", + " a = a3 ** (1 / 3)\n", "\n", + " a = a * 6.68459 * 10 ** (-12.0) # meters to AU\n", "\n", - "def a_from_Kepler3(period, M_tot):\n", - " period = period*86400 #days to seconds\n", - " \n", - "\n", - " M_tot = M_tot*M_sun.value #solar masses to kg\n", - " \n", - " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", - " \n", - " a = a3**(1/3)\n", - " \n", - " \n", - " \n", - " a = a * 6.68459*10**(-12.) # meters to AU\n", - " \n", - " return(a) #in AUs\n", - "\n", - "\n", - "\n" + " return a # in AUs" ] }, { @@ -1344,18 +1399,14 @@ "cell_type": "code", "execution_count": null, "id": "cfd9e0cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import arviz as az\n", "\n", - "az.summary(trace, var_names=[\n", - " \"P\",\n", - " \"ecc\",\n", - " \"m_planet\",\n", - " \"Omega\",\n", - " \"omega\",\n", - " \"incl\"])\n" + "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" ] }, { @@ -1372,12 +1423,14 @@ "cell_type": "code", "execution_count": null, "id": "6775ae19", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import corner\n", "\n", - "make_plots(trace, planet_params, 2)\n" + "make_plots(trace, planet_params, 2)" ] }, { @@ -1396,34 +1449,31 @@ "outputs": [], "source": [ "import pandas as pd\n", - "#save simulated data as dataframe \n", + "\n", + "# save simulated data as dataframe\n", "simulated_data_dic = {\n", " \"times_rv_observed\": x_rv,\n", " \"rv_observed\": y_rv,\n", " \"rv_err_observed\": y_rv_err,\n", - "\n", " \"times_rv_orbit\": times_rv,\n", " \"rv_orbit\": rv_orbit,\n", " \"rv_orbit_sum\": rv_orbit_sum,\n", - "\n", " \"times_astrometry_observed\": x_astrometry,\n", " \"ra_observed\": ra_data,\n", " \"dec_observed\": dec_data,\n", " \"ra_err_observed\": ra_err,\n", " \"dec_err_observed\": dec_err,\n", - "\n", " \"times_astrometry_orbit\": times_astrometry,\n", " \"ra_orbit\": ra_orbit,\n", " \"ra_orbit_sum\": ra_orbit_sum,\n", " \"dec_orbit\": dec_orbit,\n", - " \"dec_orbit_sum\": dec_orbit_sum\n", - " }\n", + " \"dec_orbit_sum\": dec_orbit_sum,\n", + "}\n", "\n", "\n", - "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient=\"index\")\n", "simulated_data = simulated_data.transpose()\n", - "simulated_data = simulated_data.fillna(value=np.nan)\n", - "\n" + "simulated_data = simulated_data.fillna(value=np.nan)" ] }, { @@ -1435,55 +1485,48 @@ "source": [ "def remove_nans(array):\n", " array = array[~pd.isnull(array)]\n", - " \n", - " return array\n", - "\n", "\n", + " return array\n", "\n", "\n", "def import_orbit_data(orbit_row):\n", " import math\n", - " \n", + "\n", " orbit_row = remove_nans(orbit_row)\n", " orbit_array = []\n", " for row in orbit_row:\n", " orbit_array.append([np.array(row[1]), row[0]])\n", "\n", " orbit_array = np.array(orbit_array)\n", - " \n", - " return orbit_array\n", - "\n", "\n", + " return orbit_array\n", "\n", "\n", "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", " import matplotlib.pyplot as plt\n", "\n", - " \n", - " \n", - " #simulated RV observations\n", - " x_rv = simulated_data['times_rv_observed'].values\n", - " y_rv = simulated_data['rv_observed'].values\n", - " y_rv_err = simulated_data['rv_err_observed'].values\n", - " \n", - " #simulated RV full orbits\n", - " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", - " y_rv_orbit = simulated_data['rv_orbit'].values\n", - " \n", - " #simulated astrometric observations\n", - " x_astrometry = simulated_data['times_astrometry_observed'].values\n", - " y_ra = simulated_data['ra_observed'].values\n", - " y_dec = simulated_data['dec_observed'].values\n", - " y_ra_err = simulated_data['ra_err_observed'].values\n", - " y_dec_err = simulated_data['dec_err_observed'].values\n", - " \n", - " #simulated astrometric full orbits\n", - " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", - " y_ra_orbit = simulated_data['ra_orbit'].values\n", - " y_dec_orbit = simulated_data['dec_orbit'].values\n", - " \n", - " \n", - " #remove any nans from simulated data (created because of mismatching lengths)\n", + " # simulated RV observations\n", + " x_rv = simulated_data[\"times_rv_observed\"].values\n", + " y_rv = simulated_data[\"rv_observed\"].values\n", + " y_rv_err = simulated_data[\"rv_err_observed\"].values\n", + "\n", + " # simulated RV full orbits\n", + " x_rv_orbit = simulated_data[\"times_rv_orbit\"].values\n", + " y_rv_orbit = simulated_data[\"rv_orbit\"].values\n", + "\n", + " # simulated astrometric observations\n", + " x_astrometry = simulated_data[\"times_astrometry_observed\"].values\n", + " y_ra = simulated_data[\"ra_observed\"].values\n", + " y_dec = simulated_data[\"dec_observed\"].values\n", + " y_ra_err = simulated_data[\"ra_err_observed\"].values\n", + " y_dec_err = simulated_data[\"dec_err_observed\"].values\n", + "\n", + " # simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data[\"times_astrometry_orbit\"].values\n", + " y_ra_orbit = simulated_data[\"ra_orbit\"].values\n", + " y_dec_orbit = simulated_data[\"dec_orbit\"].values\n", + "\n", + " # remove any nans from simulated data (created because of mismatching lengths)\n", " x_rv = remove_nans(x_rv)\n", " y_rv = remove_nans(y_rv)\n", " y_rv_err = remove_nans(y_rv_err)\n", @@ -1497,43 +1540,34 @@ " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " # add NaN for plotting purpose on true model\n", - " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", - " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", - " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", - " \n", - " \n", - " \n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", + " # add NaN for plotting purpose on true model\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float(\"NaN\"))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float(\"NaN\"), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float(\"NaN\"), axis=0)\n", "\n", - " \n", - " \n", " figsAs = []\n", " figsRV = []\n", " for n in range(0, nplanets):\n", - " \n", " ######\n", - " #make rv plots\n", - " \n", - " \n", - " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", + " # make rv plots\n", + "\n", + " figRV, ax = plt.subplots(\n", + " 2, 1, figsize=[13, 6], gridspec_kw={\"height_ratios\": [3, 1]}\n", + " )\n", " ax0, ax1 = ax[0], ax[1]\n", - " \n", + "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1544,86 +1578,68 @@ " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # Plot the data\n", - " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax0.errorbar(\n", + " x_rv,\n", + " y_rv - other,\n", + " yerr=y_rv_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_rv = []\n", " for an_x in x_rv:\n", " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", "\n", - "\n", - "\n", - " \n", " # plot the MCMC model\n", " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_rv, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "\n", + " # plot the true RA model\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color=\"k\", label=\"truth\")\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", - " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - "\n", "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " \n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_rv, y_rv - other - med_rv,\n", - " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", + " x_rv,\n", + " y_rv - other - med_rv,\n", + " yerr=y_rv_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", - " \n", + "\n", " figRV.tight_layout()\n", " figRV.subplots_adjust(hspace=0.5)\n", - " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", + " figRV.suptitle(\"planet {0}\".format(n + 1), fontsize=33)\n", " figRV.tight_layout()\n", - " \n", - " \n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", + "\n", " figRV.show()\n", " ################\n", - " #plot astrometry\n", - " \n", - " \n", - "\n", - " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", + " # plot astrometry\n", + "\n", + " figAs, ax = plt.subplots(\n", + " 2,\n", + " 2,\n", + " figsize=[18, 9],\n", + " sharey=\"row\",\n", + " gridspec_kw={\"height_ratios\": [3, 1]},\n", + " )\n", " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", "\n", " # Get the posterior median orbital parameters\n", @@ -1636,14 +1652,16 @@ " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the DEC data\n", - " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + " ax0.errorbar(\n", + " x_astrometry,\n", + " y_dec - other,\n", + " yerr=y_dec_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " # plot the MCMC DEC model\n", " pred = np.percentile(\n", @@ -1651,55 +1669,78 @@ " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_dec = []\n", " for an_x in x_astrometry:\n", " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", - " \n", - "\n", "\n", " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", "\n", " # plot the true DEC model\n", - " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - " \n", - " \n", - " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", - " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", - " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - "\n", - "\n", + " ax0.plot(\n", + " x_astrometry_orbit, y_dec_orbit.T[n], color=\"k\", label=\"truth\"\n", + " )\n", "\n", + " ax0.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax0.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax0.transAxes\n", + " )\n", + " ax0.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax0.transAxes\n", + " )\n", + " ax0.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax0.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", - " \n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - " \n", - " \n", - " #plot residuals \n", + "\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_astrometry, y_dec - other - med_dec,\n", - " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_dec - other - med_dec,\n", + " yerr=y_dec_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", "\n", - "\n", - "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1709,86 +1750,99 @@ " other = np.median(\n", " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", - " \n", - "\n", - "\n", - "\n", - "\n", "\n", " # Plot the RA data\n", - " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax2.errorbar(\n", + " x_astrometry,\n", + " y_ra - other,\n", + " yerr=y_ra_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", " [16, 50, 84],\n", - " axis=(0, 1), \n", + " axis=(0, 1),\n", " )\n", - " \n", - " \n", + "\n", " med_ra = []\n", " for an_x in x_astrometry:\n", " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", "\n", - "\n", " # plot the MCMC RA model\n", " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax2.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax2.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - " \n", - " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", - " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", - " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", "\n", + " # plot the true RA model\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color=\"k\", label=\"truth\")\n", + "\n", + " ax2.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax2.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax2.transAxes\n", + " )\n", + " ax2.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax2.transAxes\n", + " )\n", + " ax2.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax2.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", - " \n", " ax2.legend(fontsize=9, loc=1)\n", " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", - " \n", - " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax3.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax3.axhline(0.0, color=\"k\")\n", " ax3.errorbar(\n", - " x_astrometry, y_ra - other - med_ra,\n", - " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_ra - other - med_ra,\n", + " yerr=y_ra_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", + "\n", " figAs.tight_layout()\n", - " figAs.subplots_adjust(wspace=.2)\n", - " \n", - " \n", + " figAs.subplots_adjust(wspace=0.2)\n", + "\n", " figAs.show()\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "\n", - " \n", + "\n", " return None" ] }, @@ -1796,16 +1850,20 @@ "cell_type": "code", "execution_count": null, "id": "b3941ebb", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ - "make_rv_astrometry_plots(simulated_data, trace, 2)\n" + "make_rv_astrometry_plots(simulated_data, trace, 2)" ] }, { "cell_type": "markdown", "id": "1f6c90cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "### And lastly, check how well our MCMC median +/- 1 sigma results agree with the simulated values" ] @@ -1818,9 +1876,7 @@ "outputs": [], "source": [ "def load_posterior_params(trace, nplanets):\n", - "\n", - " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", - "\n", + " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\n", "\n", " params_earth = defaultdict(list)\n", " params_jup = defaultdict(list)\n", @@ -1835,138 +1891,139 @@ " posterior_meds = defaultdict(list)\n", " posterior_errs = defaultdict(list)\n", " for param in parameters:\n", - "\n", - "\n", " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", "\n", - " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", - "\n", - "\n", + " planet_quantile = [\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84),\n", + " ]\n", "\n", " planet_err = [\n", " planet_med - planet_quantile[0],\n", - " planet_quantile[1] - planet_med \n", + " planet_quantile[1] - planet_med,\n", " ]\n", "\n", - "\n", - "\n", - " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + " # print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", "\n", " posterior_meds[param] = np.round(planet_med, 3)\n", " posterior_errs[param] = np.round(planet_err, 3)\n", - " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", - " str(np.round(planet_err[1], 3)) + \" -\" + \n", - " str(np.round(planet_err[0], 3)) + \")\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " table_params[param].append(\n", + " str(np.round(planet_med, 3))\n", + " + \" (+\"\n", + " + str(np.round(planet_err[1], 3))\n", + " + \" -\"\n", + " + str(np.round(planet_err[0], 3))\n", + " + \")\"\n", + " )\n", "\n", " table_params_all.append(table_params)\n", " posterior_meds_all.append(posterior_meds)\n", " posterior_errs_all.append(posterior_errs)\n", - " #print(\"\")\n", - " #print(\"\")\n", - "\n", - "\n", + " # print(\"\")\n", + " # print(\"\")\n", "\n", " return table_params_all, posterior_meds_all, posterior_errs_all\n", "\n", "\n", - "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", - " posterior_meds, posterior_errs, nplanets):\n", - "\n", + "def make_comparison_figures(\n", + " input_params,\n", + " posterior_orbit_params,\n", + " keys_orbit_params,\n", + " posterior_meds,\n", + " posterior_errs,\n", + " nplanets,\n", + "):\n", " import matplotlib\n", " import matplotlib.pyplot as plt\n", - " \n", - " \n", + "\n", " input_params = np.array(input_params)\n", - " input_params.T[-1] = input_params.T[-1]*332946.07\n", + " input_params.T[-1] = input_params.T[-1] * 332946.07\n", "\n", " comparisons = []\n", " for ii in range(0, nplanets):\n", - " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", + " fig, axs = plt.subplots(\n", + " 1, 2, figsize=[13, 6], gridspec_kw={\"width_ratios\": [1, 1]}\n", + " )\n", " plt.subplots_adjust(hspace=1)\n", - " \n", + "\n", " #### add comparison plots\n", - " \n", - " \n", + "\n", " input_mass = input_params[ii][-1]\n", " input_period = input_params[ii][0]\n", - " \n", - " xs = posterior_meds[ii]['m_planet']\n", - " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", - "\n", - " ys = posterior_meds[ii]['P']\n", - " ys_err = np.array([posterior_errs[ii]['P']]).T\n", - "\n", "\n", + " xs = posterior_meds[ii][\"m_planet\"]\n", + " xs_err = np.array([posterior_errs[ii][\"m_planet\"]]).T\n", + "\n", + " ys = posterior_meds[ii][\"P\"]\n", + " ys_err = np.array([posterior_errs[ii][\"P\"]]).T\n", + "\n", + " color = \"#1c245f\"\n", + "\n", + " axs[0].errorbar(\n", + " xs,\n", + " ys,\n", + " xerr=xs_err,\n", + " yerr=ys_err,\n", + " marker=\"o\",\n", + " color=color,\n", + " markersize=13,\n", + " label=r\"med and 1$\\sigma$ errs\",\n", + " linewidth=3,\n", + " alpha=0.7,\n", + " )\n", "\n", - " color = '#1c245f'\n", - "\n", - " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", - " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", - " linewidth = 3, alpha = 0.7)\n", - "\n", - " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", - " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", - "\n", + " axs[0].axvline(\n", + " x=input_mass, ymin=0, ymax=1, color=\"k\", ls=\"--\", label=\"simulated\"\n", + " )\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color=\"k\", ls=\"--\")\n", "\n", - " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", - " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", + " axs[0].set_xlabel(r\"mass [M$_{earth}$]\", fontsize=27)\n", + " axs[0].set_ylabel(r\"period [days]\", fontsize=27)\n", "\n", + " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", "\n", - " fig.suptitle('input and posterior parameters', fontsize = 27)\n", + " axs[0].legend(fontsize=13)\n", "\n", - " axs[0].legend(fontsize = 13)\n", - " \n", - " \n", - " \n", " #### add comparison tables\n", "\n", - "\n", " sim_params = input_params[ii]\n", " post_params = posterior_orbit_params[ii]\n", - " data = np.column_stack((np.array(keys_orbit_params), \n", - " np.round(np.array(sim_params),3),\n", - " np.array(list(post_params.values()))))\n", + " data = np.column_stack(\n", + " (\n", + " np.array(keys_orbit_params),\n", + " np.round(np.array(sim_params), 3),\n", + " np.array(list(post_params.values())),\n", + " )\n", + " )\n", "\n", - " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", + " labels = [\"param\", \"truth\", r\"MCMC med +/- 1$\\sigma$\"]\n", "\n", " df = pd.DataFrame(data, columns=labels)\n", "\n", - " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", - " colColours =[\"#D7E5F0\"]*3, loc='center', \n", - " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", - " #for (row, col), cell in table.get_celld().items():\n", + " table = axs[1].table(\n", + " cellText=df.values,\n", + " colLabels=df.columns,\n", + " colColours=[\"#D7E5F0\"] * 3,\n", + " loc=\"center\",\n", + " bbox=[0, 0, 1, 1],\n", + " colWidths=[0.2, 0.2, 0.6],\n", + " )\n", + " # for (row, col), cell in table.get_celld().items():\n", " # if (row == 0):\n", " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", "\n", - "\n", " table.auto_set_column_width(col=labels)\n", "\n", - " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", - " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", - " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " # ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " # t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " # t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", " axs[1].set_axis_off()\n", "\n", - "\n", " table.auto_set_font_size(False)\n", " table.set_fontsize(13)\n", "\n", - " \n", - " \n", - " \n", - "\n", - "\n", - "\n", " comparisons.append(fig)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " return comparisons" ] }, @@ -1980,12 +2037,29 @@ "from collections import defaultdict\n", "\n", "\n", - "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", - "load_posterior_params(trace, 2)\n", - "\n", - "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", - "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", - " posterior_meds_all, posterior_errs_all, 2)" + "(\n", + " table_params_all,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + ") = load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = [\n", + " \"period\",\n", + " \"ecc\",\n", + " r\"T$_\\mathrm{per}$\",\n", + " r\"$\\omega$\",\n", + " r\"$\\Omega$\",\n", + " \"inc\",\n", + " \"mass\",\n", + "]\n", + "make_comparison_figures(\n", + " [planet_params[1], planet_params[0]],\n", + " table_params_all,\n", + " planet_param_keys,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + " 2,\n", + ")" ] }, { diff --git a/docs/tutorials/eb.ipynb b/docs/tutorials/eb.ipynb index 7928d814..5fbdd7e1 100644 --- a/docs/tutorials/eb.ipynb +++ b/docs/tutorials/eb.ipynb @@ -191,9 +191,7 @@ "\n", "\n", "def build_model(mask):\n", - "\n", " with pm.Model() as model:\n", - "\n", " # Systemic parameters\n", " mean_lc = pm.Normal(\"mean_lc\", mu=0.0, sd=5.0)\n", " mean_rv = pm.Normal(\"mean_rv\", mu=0.0, sd=50.0)\n", diff --git a/docs/tutorials/joint.ipynb b/docs/tutorials/joint.ipynb index 8abeff77..4ea367e5 100644 --- a/docs/tutorials/joint.ipynb +++ b/docs/tutorials/joint.ipynb @@ -311,7 +311,9 @@ "cell_type": "code", "execution_count": null, "id": "8000897c", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import pymc3 as pm\n", @@ -330,7 +332,6 @@ " if mask is None:\n", " mask = np.ones(len(x), dtype=bool)\n", " with pm.Model() as model:\n", - "\n", " # Parameters for the stellar properties\n", " mean_flux = pm.Normal(\"mean_flux\", mu=0.0, sd=10.0)\n", " u_star = xo.QuadLimbDark(\"u_star\")\n", @@ -487,7 +488,9 @@ { "cell_type": "markdown", "id": "d7d02ab9", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "Now let's plot the map radial velocity model." ] @@ -496,7 +499,9 @@ "cell_type": "code", "execution_count": null, "id": "b78425b5", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "def plot_rv_curve(soln):\n", @@ -525,7 +530,9 @@ { "cell_type": "markdown", "id": "ade1ac5b", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "That looks pretty similar to what we got in {ref}`rv`.\n", "Now let's also plot the transit model." diff --git a/docs/tutorials/lc-gp-transit.ipynb b/docs/tutorials/lc-gp-transit.ipynb index 7e3e81b4..efd1017c 100644 --- a/docs/tutorials/lc-gp-transit.ipynb +++ b/docs/tutorials/lc-gp-transit.ipynb @@ -140,12 +140,10 @@ "\n", "\n", "def build_model(mask=None, start=None):\n", - "\n", " if mask is None:\n", " mask = np.ones(len(x), dtype=bool)\n", "\n", " with pm.Model() as model:\n", - "\n", " # Shared parameters\n", " mean = pm.Normal(\"mean\", mu=0, sd=1, testval=0)\n", "\n", @@ -323,7 +321,6 @@ "\n", "\n", "def plot_light_curve(x, y, soln, mask=None):\n", - "\n", " if mask is None:\n", " mask = np.ones(len(x), dtype=bool)\n", "\n", @@ -535,7 +532,6 @@ "RUN_THE_SAMPLING = 0\n", "\n", "if RUN_THE_SAMPLING:\n", - "\n", " with model:\n", " trace = pm.sample(\n", " tune=1500,\n", diff --git a/docs/tutorials/lc-multi.ipynb b/docs/tutorials/lc-multi.ipynb index 3b79d912..d35ae4fb 100644 --- a/docs/tutorials/lc-multi.ipynb +++ b/docs/tutorials/lc-multi.ipynb @@ -171,7 +171,6 @@ "# Do several rounds of sigma clipping\n", "for i in range(10):\n", " with pm.Model() as model:\n", - "\n", " # Shared orbital parameters\n", " log_period = pm.Normal(\"log_period\", mu=np.log(lit_period), sigma=1.0)\n", " period = pm.Deterministic(\"period\", tt.exp(log_period))\n", @@ -197,7 +196,6 @@ " gp_preds = dict()\n", " gp_preds_with_mean = dict()\n", " for n, (name, (x, y, yerr, texp)) in enumerate(datasets.items()):\n", - "\n", " # We define the per-instrument parameters in a submodel so that we\n", " # don't have to prefix the names manually\n", " with pm.Model(name=name, model=model):\n", diff --git a/docs/tutorials/quick-tess.ipynb b/docs/tutorials/quick-tess.ipynb index db98cdbe..2554be3e 100644 --- a/docs/tutorials/quick-tess.ipynb +++ b/docs/tutorials/quick-tess.ipynb @@ -169,7 +169,6 @@ "\n", "\n", "with pm.Model() as model:\n", - "\n", " # Stellar parameters\n", " mean = pm.Normal(\"mean\", mu=0.0, sigma=10.0)\n", " u = xo.QuadLimbDark(\"u\")\n", diff --git a/docs/tutorials/rv-multi.ipynb b/docs/tutorials/rv-multi.ipynb index d5e2bb74..b5a9f440 100644 --- a/docs/tutorials/rv-multi.ipynb +++ b/docs/tutorials/rv-multi.ipynb @@ -120,7 +120,6 @@ "t_phase = np.linspace(-0.5, 0.5, 5000)\n", "\n", "with pm.Model() as model:\n", - "\n", " # Parameters describing the orbit\n", " log_K = pm.Normal(\"log_K\", mu=np.log(300), sigma=10)\n", " log_P = pm.Normal(\"log_P\", mu=np.log(2093.07), sigma=10)\n", diff --git a/docs/tutorials/rv.ipynb b/docs/tutorials/rv.ipynb index eb482025..5cfbceea 100644 --- a/docs/tutorials/rv.ipynb +++ b/docs/tutorials/rv.ipynb @@ -134,7 +134,6 @@ "import aesara_theano_fallback.tensor as tt\n", "\n", "with pm.Model() as model:\n", - "\n", " # Gaussian priors based on transit data (from Petigura et al.)\n", " t0 = pm.Normal(\"t0\", mu=np.array(t0s), sd=np.array(t0_errs), shape=2)\n", " logP = pm.Normal(\n", diff --git a/docs/tutorials/stellar-variability.ipynb b/docs/tutorials/stellar-variability.ipynb index 6c79ec99..9febaa6b 100644 --- a/docs/tutorials/stellar-variability.ipynb +++ b/docs/tutorials/stellar-variability.ipynb @@ -125,7 +125,6 @@ "from celerite2.theano import terms, GaussianProcess\n", "\n", "with pm.Model() as model:\n", - "\n", " # The mean flux of the time series\n", " mean = pm.Normal(\"mean\", mu=0.0, sigma=10.0)\n", "\n", diff --git a/docs/tutorials/tess.ipynb b/docs/tutorials/tess.ipynb index 34269926..7bffef0f 100644 --- a/docs/tutorials/tess.ipynb +++ b/docs/tutorials/tess.ipynb @@ -159,7 +159,9 @@ "cell_type": "code", "execution_count": null, "id": "fbb85aad", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -176,7 +178,6 @@ " if mask is None:\n", " mask = np.ones(len(x), dtype=bool)\n", " with pm.Model() as model:\n", - "\n", " # Parameters for the stellar properties\n", " mean = pm.Normal(\"mean\", mu=0.0, sd=10.0)\n", " u_star = xo.QuadLimbDark(\"u_star\")\n", @@ -307,7 +308,9 @@ { "cell_type": "markdown", "id": "82a9efd6", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "Here's how we plot the initial light curve model:" ] diff --git a/docs/tutorials/transit.ipynb b/docs/tutorials/transit.ipynb index b2596743..9f48c582 100644 --- a/docs/tutorials/transit.ipynb +++ b/docs/tutorials/transit.ipynb @@ -113,7 +113,6 @@ "import pymc3_ext as pmx\n", "\n", "with pm.Model() as model:\n", - "\n", " # The baseline flux\n", " mean = pm.Normal(\"mean\", mu=0.0, sd=1.0)\n", "\n", diff --git a/docs/tutorials/ttv.ipynb b/docs/tutorials/ttv.ipynb index 4b3298f5..25eeafa4 100644 --- a/docs/tutorials/ttv.ipynb +++ b/docs/tutorials/ttv.ipynb @@ -119,7 +119,6 @@ "np.random.seed(9485023)\n", "\n", "with pm.Model() as model:\n", - "\n", " # This part of the model is similar to the model in the `transit` tutorial\n", " mean = pm.Normal(\"mean\", mu=0.0, sd=1.0)\n", " u = xo.QuadLimbDark(\"u\", testval=np.array([0.3, 0.2]))\n", From f42897b435c2ff3707a2e7df46acdfd2b4f175d1 Mon Sep 17 00:00:00 2001 From: dyahalomi Date: Wed, 15 Mar 2023 11:27:28 -0600 Subject: [PATCH 3/6] fix heading size issue --- docs/tutorials/astrometry_rv_case_study.ipynb | 1746 ++++++++--------- 1 file changed, 836 insertions(+), 910 deletions(-) diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb index 7dcc7410..af760fdd 100644 --- a/docs/tutorials/astrometry_rv_case_study.ipynb +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -12,9 +12,7 @@ "cell_type": "code", "execution_count": null, "id": "4abd3961", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -22,7 +20,8 @@ "\n", "\n", "xo.utils.docs_setup()\n", - "print(f\"exoplanet.__version__ = '{xo.__version__}'\")" + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", + "\n" ] }, { @@ -30,7 +29,7 @@ "id": "9a87ff47", "metadata": {}, "source": [ - "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023.\n", + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", "\n", "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", "\n", @@ -50,14 +49,12 @@ { "cell_type": "markdown", "id": "93ceb2a9", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## (Simulated) Data\n", "\n", "\n", - "First, we need to simulate our dataset.\n", + "First, we need to simulate our dataset. \n", "\n", "\n", "### Let's define 2 helpful functions in order to simulate our data...." @@ -71,24 +68,27 @@ "outputs": [], "source": [ "def create_orbit(n_planets, orbit_params):\n", - " \"\"\"\n", + " \n", + " '''\n", " inputs:\n", " - n_planets = number of planets\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", " - period = period in days\n", - " - ecc = eccentricity\n", + " - ecc = eccentricity \n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node\n", + " - Omega = longitude of ascending node \n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses\n", - " \"\"\"\n", + " - m_planet = mass of planet in Earth masses \n", + " '''\n", + " \n", "\n", + " \n", " period = []\n", " ecc = []\n", - " Tper = []\n", - " omega = []\n", - " Omega = []\n", + " Tper = [] \n", + " omega = [] \n", + " Omega = [] \n", " incl = []\n", " m_planet = []\n", " for ii in range(0, n_planets):\n", @@ -101,15 +101,15 @@ " m_planet.append(orbit_params[ii][6])\n", "\n", " orbit = xo.orbits.KeplerianOrbit(\n", - " period=period,\n", - " ecc=ecc,\n", - " t_periastron=Tper,\n", - " omega=omega,\n", - " Omega=Omega,\n", - " incl=incl,\n", - " m_planet=m_planet,\n", - " )\n", - "\n", + " period = period,\n", + " ecc = ecc, \n", + " t_periastron = Tper, \n", + " omega = omega, \n", + " Omega = Omega, \n", + " incl = incl,\n", + " m_planet = m_planet) \n", + " \n", + " \n", " return orbit" ] }, @@ -120,24 +120,24 @@ "metadata": {}, "outputs": [], "source": [ - "np.random.seed(979)\n", - "\n", + "np.random.seed(979) \n", "\n", "def simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra,\n", " sigma_dec,\n", " plx,\n", " orbit_params,\n", - " times_observed_rv=None,\n", - " times_observed_astrometry=None,\n", - " t_dur_rv=None,\n", - " n_obs_rv=None,\n", - " t_dur_astrometry=None,\n", - " n_obs_astrometry=None,\n", - "):\n", - " \"\"\"\n", + " times_observed_rv = None,\n", + " times_observed_astrometry = None,\n", + " t_dur_rv = None, \n", + " n_obs_rv = None,\n", + " t_dur_astrometry = None, \n", + " n_obs_astrometry = None \n", + " ):\n", + "\n", + " '''\n", " inputs:\n", " - orbit = keplerian orbit generated by exoplanet code\n", " - n_planets = number of planets\n", @@ -147,231 +147,230 @@ " - plx = parallax of system\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", " - period = period in days\n", - " - ecc = eccentricity\n", + " - ecc = eccentricity \n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node\n", + " - Omega = longitude of ascending node \n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses\n", + " - m_planet = mass of planet in Earth masses \n", " - times_observed_rv = (default=None) the observed times for RV observations\n", " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", - "\n", - "\n", - " \"\"\"\n", - "\n", + " \n", + " \n", + " '''\n", + " \n", " if n_planets == 1:\n", " orbit_params = [orbit_params]\n", - "\n", - " # if times_observed_rv != None then just define times_rv and linspace between min and max\n", - " if times_observed_rv != None:\n", - " times_rv = np.linspace(\n", - " np.min(times_observed_rv), np.max(times_observed_rv), 10000\n", - " )\n", - "\n", - " # define start_time_RV, which will be used to determine the times to simulate\n", - " # if times_observed_RV and times_observed_astrometry aren't defined\n", + " \n", + " \n", + " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", + "\n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_rv = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_rv:\n", " start_time_rv = orbit_params[ii][2]\n", "\n", - " # define RV observation time\n", - " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_rv = np.linspace(start_time_rv, start_time_rv + t_dur_rv, 10000)\n", - " times_observed_rv = np.linspace(\n", - " start_time_rv, start_time_rv + t_dur_rv, n_obs_rv\n", - " )\n", "\n", - " # if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", - " if times_observed_astrometry != None:\n", - " times_astrometry = np.linspace(\n", - " np.min(times_observed_astrometry),\n", - " np.max(times_observed_astrometry),\n", - " 10000,\n", - " )\n", + " \n", + " #define RV observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", + " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", + " \n", + "\n", "\n", - " # define start_time_RV, which will be used to determine the times to simulate\n", - " # if times_observed_RV and times_observed_astrometry aren't defined\n", + " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", + " \n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined \n", " else:\n", " start_time_astrometry = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_astrometry:\n", - " start_time_astrometry = orbit_params[ii][2]\n", - "\n", - " # define astrometry observation time\n", - " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_astrometry = np.linspace(\n", - " start_time_astrometry,\n", - " start_time_astrometry + t_dur_astrometry,\n", - " 10000,\n", - " )\n", - " times_observed_astrometry = np.linspace(\n", - " start_time_astrometry,\n", - " start_time_astrometry + t_dur_astrometry,\n", - " n_obs_astrometry,\n", - " )\n", + " start_time_astrometry = orbit_params[ii][2] \n", "\n", - " orbit = create_orbit(n_planets, orbit_params)\n", "\n", + " #define astrometry observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", + " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", + "\n", + "\n", + " \n", + " orbit = create_orbit(n_planets, orbit_params)\n", + " \n", " rv = orbit.get_radial_velocity(times_rv)\n", " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", - "\n", + " \n", " rv_orbit = rv.eval()\n", "\n", " if n_planets > 1:\n", - " rv_orbit_sum = np.sum(rv_orbit, axis=1)\n", + " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", " else:\n", " rv_orbit_sum = rv_orbit\n", + " \n", "\n", " rv_observed = rv_observed.eval()\n", " if n_planets > 1:\n", - " rv_observed_sum = np.sum(rv_observed, axis=1)\n", + " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", " else:\n", " rv_observed_sum = rv_observed\n", "\n", - " # -----------\n", - " # -----------\n", + "\n", + " #-----------\n", + " #-----------\n", "\n", " # determine and print the star position at desired times\n", " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", "\n", - " # pos = tt.sum(pos, axis=-1)\n", "\n", - " x, y, z = pos\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x,y,z = pos\n", + "\n", "\n", " # calculate rho and theta\n", - " rho = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", + " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - "\n", + " \n", " rho, theta = rho.eval(), theta.eval()\n", + " \n", + " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", - " # rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", " rho_orbit = rho\n", - "\n", + " \n", " if n_planets > 1:\n", - " rho_orbit_sum = np.sum(rho_orbit, axis=1)\n", + " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", " else:\n", " rho_orbit_sum = rho_orbit\n", "\n", " theta_orbit = theta\n", - "\n", + " \n", " if n_planets > 1:\n", - " theta_orbit_sum = np.sum(theta_orbit, axis=1)\n", + " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", " else:\n", " theta_orbit_sum = theta_orbit\n", - "\n", + " \n", + " \n", " # when summing over theta, position angle, we have to careful because position\n", - " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract\n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", " # in the correct range. Be careful though if modeling more than 2 planets, this\n", " # doesn't completely solve the problem!\n", - " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2 * np.pi\n", - " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2 * np.pi\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", + "\n", "\n", " # determine and print the star position at desired times\n", - " pos_observed = theano.function(\n", - " [], orbit.get_star_position(times_observed_astrometry, plx)\n", - " )()\n", + " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", "\n", - " # pos = tt.sum(pos, axis=-1)\n", "\n", - " x_obs, y_obs, z_obs = pos_observed\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x_obs,y_obs,z_obs = pos_observed\n", "\n", - " # calculate rho and theta\n", - " rho_observed = tt.squeeze(tt.sqrt(x_obs**2 + y_obs**2)) # arcsec\n", - " theta_observed = tt.squeeze(\n", - " tt.arctan2(y_obs, x_obs)\n", - " ) # radians between [-pi, pi]\n", "\n", + " # calculate rho and theta\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", + " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", + " \n", " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", "\n", + "\n", + " \n", " if n_planets > 1:\n", - " rho_observed_sum = np.sum(rho_observed, axis=1)\n", + " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", " else:\n", " rho_observed_sum = rho_observed\n", "\n", + " \n", " if n_planets > 1:\n", - " theta_observed_sum = np.sum(theta_observed, axis=1)\n", + " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", " else:\n", " theta_observed_sum = theta_observed\n", - "\n", + " \n", + " \n", " #######\n", " #######\n", " #######\n", "\n", " # convert rho and theta into ra and dec\n", - " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", "\n", " if n_planets > 1:\n", - " ra_orbit_sum = np.sum(ra_orbit, axis=1)\n", + " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", " else:\n", " ra_orbit_sum = ra_orbit\n", "\n", + " \n", " if n_planets > 1:\n", - " dec_orbit_sum = np.sum(dec_orbit, axis=1)\n", + " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", " else:\n", " dec_orbit_sum = dec_orbit\n", "\n", - " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", "\n", - " # -----------\n", - " # -----------\n", - " # ----------- simulate data\n", - " # -----------\n", - " # -----------\n", + "\n", + "\n", + "\n", + " \n", + " #----------- \n", + " #----------- \n", + " #----------- simulate data\n", + " #----------- \n", + " #----------- \n", + "\n", "\n", " if n_planets > 1:\n", " rv_sim = rv_observed\n", - " rv_sim_sum = np.sum(rv_observed, axis=1) + np.random.normal(\n", - " 0, sigma_rv, len(rv_observed)\n", - " )\n", + " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", " else:\n", " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", " rv_sim_sum = rv_sim\n", "\n", - " # -----------\n", - " # -----------\n", + " #-----------\n", + " #-----------\n", "\n", - " if n_planets > 1:\n", - " ra_sim = ra_observed\n", - " ra_sim_sum = np.sum(ra_sim, axis=1) + np.random.normal(\n", - " 0, sigma_ra, len(ra_observed)\n", - " )\n", "\n", + " if n_planets > 1:\n", + " ra_sim = ra_observed \n", + " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", + " \n", " else:\n", " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", " ra_sim_sum = ra_sim\n", - "\n", + " \n", + " \n", + " \n", " if n_planets > 1:\n", - " dec_sim = dec_observed\n", - " dec_sim_sum = np.sum(dec_sim, axis=1) + np.random.normal(\n", - " 0, sigma_dec, len(dec_observed)\n", - " )\n", - "\n", + " dec_sim = dec_observed \n", + " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " \n", " else:\n", - " dec_sim = dec_observed + np.random.normal(\n", - " 0, sigma_dec, len(dec_observed)\n", - " )\n", + " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", " dec_sim_sum = dec_sim\n", "\n", - " times = [\n", - " times_rv,\n", - " times_observed_rv,\n", - " times_astrometry,\n", - " times_observed_astrometry,\n", - " ]\n", + " \n", + " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", - "\n", + " \n", + " \n", " return [times, rv_results, ra_results, dec_results]" ] }, @@ -399,20 +398,10 @@ "source": [ "import numpy as np\n", "\n", - "parallax = 0.1 # as\n", - "n_planets = 2\n", - "planet_params = [\n", - " [365.0, 0.0167, 100.0, 1.7959438003021653, 0.0, np.pi / 4, 3e-06],\n", - " [\n", - " 4327.0,\n", - " 0.0484,\n", - " 500.0,\n", - " -1.4957471689591395,\n", - " 1.7523105690023069,\n", - " np.pi / 4 + 0.022759,\n", - " 0.000955,\n", - " ],\n", - "]" + "parallax = 0.1 #as \n", + "n_planets=2\n", + "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", + " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" ] }, { @@ -433,28 +422,28 @@ "cell_type": "code", "execution_count": null, "id": "55938572", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "from aesara_theano_fallback import aesara as theano\n", "import theano.tensor as tt\n", "\n", "\n", - "# add gaia observing times\n", + "#add gaia observing times\n", "times_observed_astrometry_gaia = []\n", "t_0 = 0\n", - "for ii in range(t_0, t_0 + 3600):\n", - " if ii % (3600 / 100) == 0:\n", + "for ii in range(t_0, t_0+3600):\n", + " if ii % (3600/100) == 0:\n", " times_observed_astrometry_gaia.append(ii)\n", "\n", "\n", - "# add THE observing times\n", + "\n", + "#add THE observing times\n", "times_observed_rv = []\n", - "t_0 = 3300 # Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", "add_data = True\n", - "for ii in range(t_0, t_0 + 3600):\n", + "for ii in range(t_0, t_0+3600):\n", + "\n", " if ii % 180 == 0:\n", " if add_data:\n", " add_data = False\n", @@ -465,66 +454,72 @@ " times_observed_rv.append(ii)\n", "\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "sigma_rv = 0.3\n", "\n", - "sigma_ra_gaia = 6e-5 # 60 micro-as\n", - "sigma_dec_gaia = 6e-5 # 60 micro-as\n", + "sigma_ra_gaia = 6e-5 #60 micro-as\n", + "sigma_dec_gaia = 6e-5 #60 micro-as\n", + "\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra_gaia,\n", " sigma_dec_gaia,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv=times_observed_rv,\n", - " times_observed_astrometry=times_observed_astrometry_gaia,\n", - ")\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_gaia\n", + " )\n", "\n", "\n", - "[\n", - " [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", - " [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", - " [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", - " [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum],\n", - "] = (times, rv_results, ra_results, dec_results)\n", + "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", "\n", "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", "\n", "\n", + "\n", + "\n", + "\n", "# we assume that the first roman observation\n", "# will occur ~5 years after the last Gaia observation\n", - "t_1 = times_observed_astrometry_gaia[-1] + 1800\n", + "t_1 = times_observed_astrometry_gaia[-1]+1800\n", "\n", "# add roman observing times\n", "times_observed_astrometry_roman = []\n", - "for ii in range(t_1, t_1 + (10 * 365)):\n", + "for ii in range(t_1, t_1+(10*365)):\n", " if ii % 90 == 0:\n", - " times_observed_astrometry_roman.append(ii)\n", + " times_observed_astrometry_roman.append(ii) \n", + "\n", "\n", + "sigma_ra_roman = 5e-6 #5 micro-as\n", + "sigma_dec_roman = 5e-6 #5 micro-as\n", "\n", - "sigma_ra_roman = 5e-6 # 5 micro-as\n", - "sigma_dec_roman = 5e-6 # 5 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra_roman,\n", " sigma_dec_roman,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv=times_observed_rv,\n", - " times_observed_astrometry=times_observed_astrometry_roman,\n", - ")\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_roman\n", + " )\n", "\n", "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", "\n", - "times_observed_astrometry = np.append(\n", - " times_observed_astrometry, times[3], axis=0\n", - ")\n", + "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", "\n", "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", @@ -537,7 +532,7 @@ "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", "\n", "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", - "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)" + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" ] }, { @@ -557,40 +552,36 @@ "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib\n", - "\n", - "matplotlib.rc(\"font\", **{\"family\": \"serif\", \"serif\": [\"Computer Modern\"]})\n", - "matplotlib.rc(\"text\", usetex=True)\n", - "\n", + "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", + "matplotlib.rc('text', usetex=True)\n", "\n", "def plot_rv_signal(\n", - " n_planets,\n", - " rv_orbit,\n", - " rv_orbit_sum,\n", - " rv_sim,\n", - " rv_sim_sum,\n", - " times_rv,\n", - " times_observed_rv,\n", - "):\n", - " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", - " fig.suptitle(\"RV Signal\", fontsize=45)\n", + " n_planets, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv):\n", + " \n", + " \n", + " \n", + " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", + " fig.suptitle(\"RV Signal\", fontsize = 45)\n", " ax1 = ax\n", + " \n", + " \n", + " \n", "\n", - " ax1.plot(\n", - " times_observed_rv,\n", - " rv_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " alpha=0.3,\n", - " )\n", + " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", "\n", - " ax1.set_xlabel(\"time [BJD]\", fontsize=27)\n", - " ax1.set_ylabel(\"RV [m/s]\", fontsize=27)\n", - " ax1.legend(fontsize=27, loc=2)\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", + " ax1.legend(fontsize = 27, loc = 2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", - "\n", + " \n", " return None" ] }, @@ -602,55 +593,53 @@ "outputs": [], "source": [ "def plot_astrometry_signal(\n", - " n_planets,\n", - " ra_orbit,\n", - " ra_orbit_sum,\n", - " ra_sim,\n", - " ra_sim_sum,\n", - " dec_orbit,\n", - " dec_orbit_sum,\n", - " dec_sim,\n", + " n_planets, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry,\n", - "):\n", - " fig, ax = plt.subplots(2, 1, figsize=[13, 10], sharex=\"col\", sharey=\"row\")\n", - " fig.suptitle(\"Astrometric Signal\", fontsize=45)\n", + " times_observed_astrometry):\n", + " \n", + " \n", + "\n", + " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", + " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", " ax2 = ax[0]\n", " ax3 = ax[1]\n", + " \n", "\n", - " ax2.plot(\n", - " times_observed_astrometry,\n", - " ra_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " )\n", + "\n", + " \n", + " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", " for tick in ax3.get_xticklabels():\n", " tick.set_rotation(30)\n", - " ax2.legend(fontsize=27, loc=2)\n", - "\n", - " ax3.plot(\n", - " times_observed_astrometry,\n", - " dec_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " )\n", - " ax2.set_xlabel(\"time [BJD]\", fontsize=27)\n", - " ax3.legend(fontsize=27, loc=2)\n", + " ax2.legend(fontsize = 27, loc = 2)\n", + " \n", + " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax3.legend(fontsize = 27, loc = 2)\n", + "\n", + "\n", "\n", " fig.tight_layout()\n", " fig.show()\n", "\n", + "\n", + "\n", " #######\n", " #######\n", " #######\n", "\n", - " fig, ax = plt.subplots(1, 1, figsize=[10, 10])\n", + "\n", + " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", "\n", " # plot simulated data\n", - " ax.plot(ra_sim_sum, dec_sim_sum, \"o\", color=\"k\", label=\"combined signal\")\n", + " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", "\n", " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", @@ -660,7 +649,7 @@ " ax.set_title(\"initial orbit\")\n", " ax.legend()\n", " fig.show()\n", - "\n", + " \n", " return None" ] }, @@ -676,46 +665,46 @@ "cell_type": "code", "execution_count": null, "id": "82e18e64", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "plot_rv_signal(\n", - " 2, rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum, times_rv, times_observed_rv\n", - ")" + " 2, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv)\n", + " \n", + " \n" ] }, { "cell_type": "code", "execution_count": null, "id": "623b9c2d", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "plot_astrometry_signal(\n", - " 2,\n", - " ra_orbit,\n", - " ra_orbit_sum,\n", - " ra_sim,\n", - " ra_sim_sum,\n", - " dec_orbit,\n", - " dec_orbit_sum,\n", - " dec_sim,\n", + " 2, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry,\n", - ")" + " times_observed_astrometry)" ] }, { "cell_type": "markdown", "id": "99110cc0", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## Ok, now we have a simulated dataset -- let's model it!\n", "\n", @@ -736,7 +725,9 @@ " print(\"------------\")\n", "\n", " with pm.Model() as model:\n", - " # log normal prior on period around estimates\n", + " \n", + " \n", + " #log normal prior on period around estimates\n", " logP = pm.Uniform(\n", " \"logP\",\n", " lower=0,\n", @@ -744,33 +735,37 @@ " shape=2,\n", " testval=np.log(periods),\n", " )\n", + " \n", + "\n", "\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", "\n", + "\n", " ## wide uniform prior on t_periastron\n", " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", - "\n", + " \n", + " \n", " # Wide normal prior for semi-amplitude\n", - " logK = pm.Uniform(\n", - " \"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks)\n", - " )\n", - "\n", + " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", + " \n", " K = pm.Deterministic(\"K\", tt.exp(logK))\n", - "\n", + " \n", + " \n", " # Eccentricity & argument of periasteron\n", " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", "\n", + "\n", + "\n", + " \n", " # Jitter & a quadratic RV trend\n", - " # logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", "\n", " # Then we define the orbit\n", - " orbit = xo.orbits.KeplerianOrbit(\n", - " period=P, t_periastron=tperi, ecc=ecc, omega=omega\n", - " )\n", + " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", @@ -781,6 +776,8 @@ " # Sum over planets and add the background to get the full model\n", " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", + " \n", + "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", "\n", @@ -789,17 +786,19 @@ "\n", " # Finally add in the observation model. This next line adds a new contribution\n", " # to the log probability of the PyMC3 model\n", - " # err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", " err = y_rv_err\n", " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", "\n", + "\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", " map_soln = pmx.optimize(start=map_soln)\n", "\n", - " # return the max a-posteriori solution\n", + "\n", + " #return the max a-posteriori solution\n", " return map_soln" ] }, @@ -807,18 +806,16 @@ "cell_type": "code", "execution_count": null, "id": "4d6faa40", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import pymc3 as pm\n", "import pymc3_ext as pmx\n", "\n", - "np.random.seed(979)\n", + "np.random.seed(979) \n", "\n", "\n", - "# rename variables in more consistent way for modeling\n", + "#rename variables in more consistent way for modeling\n", "x_rv = np.array(times_observed_rv)\n", "y_rv = rv_sim_sum\n", "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", @@ -832,29 +829,25 @@ "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", "\n", "\n", + " \n", + " \n", "# make a fine grid that spans the observation window for plotting purposes\n", - "t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - ")\n", + "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", "# for predicted orbits\n", - "t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - ")\n", + "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", "\n", "\n", - "# minimize on RV data\n", + "#minimize on RV data\n", "periods_real = [4327.0, 365.0]\n", "periods_guess = []\n", "\n", - "# let's assume that THE can recover the period within 10% of the true value\n", + "#let's assume that THE can recover the period within 10% of the true value\n", "for per in periods_real:\n", - " periods_guess.append(\n", - " per + np.random.uniform(low=-0.1 * per, high=0.1 * per)\n", - " )\n", + " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", "\n", - "print(\"periods used as starting point = \", periods_guess)\n", + "print('periods used as starting point = ', periods_guess)\n", "periods_guess = np.array(periods_guess)\n", "\n", "\n", @@ -866,9 +859,7 @@ { "cell_type": "markdown", "id": "efc4dcc2", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## Now let's define out joint astrometry and radial velocity model:" ] @@ -880,40 +871,25 @@ "metadata": {}, "outputs": [], "source": [ - "def minimize_both(\n", - " rv_map_soln,\n", - " x_rv,\n", - " y_rv,\n", - " y_rv_err,\n", - " x_astrometry,\n", - " ra_data,\n", - " ra_err,\n", - " dec_data,\n", - " dec_err,\n", - " parallax,\n", - "):\n", - " m_sun = 333030 # earth masses\n", - "\n", - " P_RV = np.array(rv_map_soln[\"P\"])\n", - " K_RV = np.array(rv_map_soln[\"K\"])\n", - " tperi_RV = np.array(rv_map_soln[\"tperi\"])\n", - " ecc_RV = np.array(rv_map_soln[\"ecc\"])\n", - " omega_RV = np.array(rv_map_soln[\"omega\"])\n", - " min_masses_RV = (\n", - " xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value * 317.83\n", - " ) # in m_earth\n", + "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", + " m_sun = 333030 #earth masses\n", + " \n", + " P_RV = np.array(rv_map_soln['P'])\n", + " K_RV = np.array(rv_map_soln['K'])\n", + " tperi_RV = np.array(rv_map_soln['tperi'])\n", + " ecc_RV = np.array(rv_map_soln['ecc'])\n", + " omega_RV = np.array(rv_map_soln['omega'])\n", + " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", " phase_RV = determine_phase(P_RV, tperi_RV)\n", - "\n", + " \n", + " \n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - " )\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "\n", "\n", " print(\"RV Solutions\")\n", " print(\"------------\")\n", @@ -923,106 +899,98 @@ " print(\"eccentricity: \", ecc_RV)\n", " print(\"omega: \", omega_RV)\n", "\n", - " print(\"\")\n", + " print('')\n", " print(\"minimizing joint model solutions pre-MCMC\")\n", " print(\"------------\")\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", "\n", - " inc_test_vals = np.array(np.radians([5.0, 25.0, 45.0, 65.0, 85.0]))\n", + " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", " model, map_soln, logp = [], [], []\n", " for inc in inc_test_vals:\n", - " mass_test_vals = min_masses_RV / np.sin(inc)\n", - " print(\"\")\n", + " mass_test_vals = min_masses_RV/np.sin(inc)\n", + " print('')\n", " print(\"trying inclination = \" + str(np.degrees(inc)))\n", " print(\"mass test val = \" + str(mass_test_vals))\n", " print(\"------------\")\n", + " \n", + " \n", + "\n", + "\n", "\n", " def get_model():\n", " with pm.Model() as model:\n", + "\n", + "\n", " # Below we will run a version of this model where a measurement of parallax is provided\n", " # The measurement is in milliarcsec\n", " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", - " \"m_plx\", mu=parallax * 1000, sd=10, testval=parallax * 1000\n", + " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", " )\n", " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", "\n", + "\n", " # We expect the period to be around that found from just the RVs,\n", " # so we'll set a broad prior on logP\n", - "\n", + " \n", " logP = pm.Uniform(\n", - " \"logP\",\n", - " lower=0,\n", - " upper=np.log(2 * P_RV),\n", - " testval=np.log(P_RV),\n", - " shape=2,\n", + " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", " )\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", - "\n", + " \n", " # Eccentricity & argument of periasteron\n", - " ecs = pmx.UnitDisk(\n", - " \"ecs\",\n", - " shape=(2, 2),\n", - " testval=np.array(\n", - " [\n", - " np.sqrt(ecc_RV) * np.cos(omega_RV),\n", - " np.sqrt(ecc_RV) * np.sin(omega_RV),\n", - " ]\n", - " ),\n", - " )\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", + " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", + " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", + " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", - " xo.eccentricity.kipping13(\n", - " \"ecc_prior\", fixed=True, observed=ecc\n", - " )\n", + " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", + "\n", + " \n", + " \n", "\n", " # Omegas are co-dependent, so sample them with variables Omega_plus\n", - " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and\n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", " # Omega_minus is (Omega_0 - Omega_1)/2\n", - "\n", + " \n", " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", + " \n", + " \n", + " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", + " (Omega_plus - Omega_minus)] )\n", + " \n", "\n", - " Omega = tt.concatenate(\n", - " [(Omega_plus + Omega_minus), (Omega_plus - Omega_minus)]\n", - " )\n", + " Omega = pm.Deterministic(\"Omega\", Omega) \n", + " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", + " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", + " \n", "\n", - " Omega = pm.Deterministic(\"Omega\", Omega)\n", - " Omega_sum = pm.Deterministic(\n", - " \"Omega_sum\", ((Omega_plus) * 2) % np.pi\n", - " )\n", - " Omega_diff = pm.Deterministic(\n", - " \"Omega_diff\", ((Omega_minus) * 2) % np.pi\n", - " )\n", "\n", + " \n", " # For these orbits, it can also be better to fit for a phase angle\n", " # (relative to a reference time) instead of the time of periasteron\n", " # passage directly\n", " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", + " \n", "\n", + " \n", " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", - " log_m = pm.Uniform(\n", - " \"log_m\",\n", - " lower=-1,\n", - " upper=np.log(100 * mass_test_vals),\n", - " testval=np.log(mass_test_vals),\n", - " shape=2,\n", - " )\n", + " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", - " m_planet_fit = pm.Deterministic(\n", - " \"m_planet_fit\", m_planet / m_sun\n", - " )\n", + " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", "\n", - " cos_incl = pm.Uniform(\n", - " \"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2\n", - " )\n", + "\n", + " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", + " \n", + "\n", "\n", + " \n", + " \n", " # Set up the orbit\n", " orbit = xo.orbits.KeplerianOrbit(\n", " t_periastron=tperi,\n", @@ -1031,10 +999,13 @@ " ecc=ecc,\n", " omega=omega,\n", " Omega=Omega,\n", - " m_planet=m_planet_fit,\n", - " plx=plx,\n", + " m_planet = m_planet_fit,\n", + " plx=plx\n", " )\n", "\n", + "\n", + " \n", + " \n", " # Add a function for computing the full astrometry model\n", " def get_astrometry_model(t, name=\"\"):\n", " # First the astrometry induced by the planets\n", @@ -1042,56 +1013,54 @@ " # determine and print the star position at desired times\n", " pos = orbit.get_star_position(t, plx)\n", "\n", - " x, y, z = pos\n", - "\n", - " # calculate rho and theta\n", - " rhos = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", - " thetas = tt.squeeze(\n", - " tt.arctan2(y, x)\n", - " ) # radians between [-pi, pi]\n", + " x,y,z = pos\n", "\n", - " # rhos, thetas = get_star_relative_angles(t, plx)\n", - "\n", - " dec = pm.Deterministic(\n", - " \"dec\" + name, rhos * np.cos(thetas)\n", - " ) # X is north\n", - " ra = pm.Deterministic(\n", - " \"ra\" + name, rhos * np.sin(thetas)\n", - " ) # Y is east\n", "\n", + " # calculate rho and theta\n", + " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", + " \n", + " \n", + " #rhos, thetas = get_star_relative_angles(t, plx)\n", + " \n", + " \n", + " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", + " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", + " \n", " # Sum over planets to get the full model\n", - " dec_model = pm.Deterministic(\n", - " \"dec_model\" + name, tt.sum(dec, axis=-1)\n", - " )\n", - " ra_model = pm.Deterministic(\n", - " \"ra_model\" + name, tt.sum(ra, axis=-1)\n", - " )\n", + " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", + " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", + " \n", "\n", + " \n", " return dec_model, ra_model\n", "\n", + " \n", " # Define the astrometry model at the observed times\n", " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", "\n", " # Also define the model on a fine grid as computed above (for plotting)\n", - " dec_model_fine, ra_model_fine = get_astrometry_model(\n", - " t_fine, name=\"_fine\"\n", - " )\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", + "\n", + " \n", + "\n", "\n", - " # dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", - " # ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", " dec_tot_err = dec_err\n", " ra_tot_err = ra_err\n", "\n", - " # define the likelihood function, e.g., a Gaussian on both ra and dec\n", - " pm.Normal(\n", - " \"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data\n", - " )\n", - " pm.Normal(\n", - " \"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data\n", - " )\n", "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", + " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", + " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", + "\n", + "\n", + " \n", + " \n", " # ADD RV MODEL\n", "\n", + "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", " # First the RVs induced by the planets\n", @@ -1099,9 +1068,7 @@ " pm.Deterministic(\"vrad\" + name, vrad)\n", "\n", " # Sum over planets to get the full model\n", - " return pm.Deterministic(\n", - " \"rv_model\" + name, tt.sum(vrad, axis=-1)\n", - " )\n", + " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", @@ -1117,20 +1084,24 @@ " # Optimize to find the initial parameters\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", - " # map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", "\n", " map_soln = pmx.optimize(map_soln)\n", "\n", + "\n", " return model, map_soln\n", "\n", + "\n", + "\n", " a_model, a_map_soln = get_model()\n", - " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis=0)\n", - " print(\"log likelihood = \" + str(a_logp))\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", + " print('log likelihood = ' + str(a_logp))\n", "\n", " model.append(a_model)\n", " map_soln.append(a_map_soln)\n", " logp.append(a_logp)\n", "\n", + " \n", " best_index = 0\n", " for index in range(0, len(model)):\n", " if logp[index] >= logp[best_index]:\n", @@ -1140,26 +1111,31 @@ " the_map_soln = map_soln[best_index]\n", " the_logp = logp[best_index]\n", "\n", + "\n", + "\n", " return the_model, the_map_soln, the_logp\n", "\n", "\n", + "\n", "def determine_phase(P, t_periastron):\n", " phase = (2 * np.pi * t_periastron) / P\n", " return phase\n", "\n", "\n", + "\n", "def model_both(model, map_soln, tune_steps, draw_steps):\n", - " print(\"Joint RV + Astometry Minimization Solutions:\")\n", + " print('Joint RV + Astometry Minimization Solutions:')\n", " print(\"------------\")\n", "\n", - " print(\"m_planet: \" + str(map_soln[\"m_planet\"]))\n", - " print(\"P: \" + str(map_soln[\"P\"]))\n", - " print(\"incl: \" + str(map_soln[\"incl\"]))\n", - " print(\"Omega: \" + str(map_soln[\"Omega\"]))\n", - " print(\"tperi: \" + str(map_soln[\"tperi\"]))\n", - " print(\"ecc: \" + str(map_soln[\"ecc\"]))\n", - " print(\"omega: \" + str(map_soln[\"omega\"]))\n", - " print(\"plx: \" + str(map_soln[\"plx\"]))\n", + " print('m_planet: ' + str(map_soln['m_planet']))\n", + " print('P: ' + str(map_soln['P']))\n", + " print('incl: ' + str(map_soln['incl']))\n", + " print('Omega: ' + str(map_soln['Omega']))\n", + " print('tperi: ' + str(map_soln['tperi']))\n", + " print('ecc: ' + str(map_soln['ecc']))\n", + " print('omega: ' + str(map_soln['omega']))\n", + " print('plx: ' + str(map_soln['plx']))\n", + "\n", "\n", " with model:\n", " trace = pmx.sample(\n", @@ -1187,24 +1163,18 @@ "\n", "################\n", "################\n", - "# minimize on joint model\n", + "#minimize on joint model\n", "joint_model, joint_map_soln, joint_logp = minimize_both(\n", - " rv_map_soln,\n", - " x_rv,\n", - " y_rv,\n", - " y_rv_err,\n", - " x_astrometry,\n", - " ra_data,\n", - " ra_err,\n", - " dec_data,\n", - " dec_err,\n", - " parallax,\n", + " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", + " ra_data, ra_err, dec_data, dec_err, parallax\n", ")\n", "\n", "\n", + "\n", + "\n", "################\n", "################\n", - "# run full MCMC\n", + "#run full MCMC\n", "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" ] }, @@ -1223,168 +1193,143 @@ "metadata": {}, "outputs": [], "source": [ - "matplotlib.rc(\"text\", usetex=False)\n", - "\n", + "matplotlib.rc('text', usetex=False)\n", "\n", "def make_plots(trace, orbit_params, n_planets):\n", - " m_sun = 333030 # earth masses\n", + " m_sun = 333030 #earth masses\n", "\n", " [orbit_params_earth, orbit_params_jup] = orbit_params\n", "\n", - " [\n", - " P_earth,\n", - " e_earth,\n", - " Tper_earth,\n", - " omega_earth,\n", - " Omega_earth,\n", - " inclination_earth,\n", - " m_earth,\n", - " ] = orbit_params_earth\n", - "\n", - " [\n", - " P_jup,\n", - " e_jup,\n", - " Tper_jup,\n", - " omega_jup,\n", - " Omega_jup,\n", - " inclination_jup,\n", - " m_jup,\n", - " ] = orbit_params_jup\n", - "\n", - " a_true_earth = a_from_Kepler3(P_earth, 1.0 + m_earth)\n", - " a_true_jup = a_from_Kepler3(P_jup, 1.0 + m_jup)\n", - "\n", - " # [P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", - " truths = [\n", - " P_jup,\n", - " P_earth,\n", - " e_jup,\n", - " e_earth,\n", - " omega_jup,\n", - " omega_earth,\n", - " Omega_earth - Omega_jup,\n", - " Omega_earth + Omega_jup,\n", - " inclination_jup,\n", - " inclination_earth,\n", - " 0.1,\n", - " m_jup * m_sun,\n", - " m_earth * m_sun,\n", - " Tper_jup,\n", - " Tper_earth,\n", - " ]\n", - "\n", - " # [P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", - " truth_chain_plot = [\n", - " P_jup,\n", - " P_earth,\n", - " e_jup,\n", - " e_earth,\n", - " omega_jup,\n", - " omega_earth,\n", - " inclination_jup,\n", - " inclination_earth,\n", - " Tper_jup,\n", - " Tper_earth,\n", - " m_jup * m_sun,\n", - " m_earth * m_sun,\n", - " ]\n", + "\n", + " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", + " inclination_earth, m_earth] = orbit_params_earth\n", + "\n", + "\n", + " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", + " inclination_jup, m_jup] = orbit_params_jup\n", + "\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", + " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", + " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", + "\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", + " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " # plot the table summarizing MCMC results\n", " az.summary(\n", " trace,\n", - " var_names=[\n", - " \"P\",\n", - " \"tperi\",\n", - " \"omega\",\n", - " \"Omega_sum\",\n", - " \"Omega_plus\",\n", - " \"Omega_minus\",\n", - " \"incl\",\n", - " \"ecc\",\n", - " \"plx\",\n", - " \"m_planet\",\n", - " ],\n", + " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", + " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", " )\n", "\n", + "\n", " plt.show()\n", "\n", + "\n", + "\n", + "\n", + "\n", " # plot the corner plots\n", " _ = corner.corner(\n", - " trace,\n", - " var_names=[\n", - " \"P\",\n", - " \"ecc\",\n", - " \"omega\",\n", - " \"Omega_sum\",\n", - " \"Omega_minus\",\n", - " \"incl\",\n", - " \"plx\",\n", - " \"m_planet\",\n", - " \"tperi\",\n", - " ],\n", - " quantiles=[0.16, 0.5, 0.84],\n", - " show_titles=True,\n", - " title_kwargs={\"fontsize\": 13},\n", - " truths=truths,\n", - " truth_color=\"#03003a\",\n", + " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", + " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", + " truths = truths, truth_color = \"#03003a\"\n", " )\n", "\n", " plt.show()\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # plot the chains\n", " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", " for ii in range(0, len(parameters)):\n", " plot_truth = False\n", " param = parameters[ii]\n", - "\n", - " true_vals_earth = truth_chain_plot[2 * ii]\n", - " true_vals_jup = truth_chain_plot[2 * ii + 1]\n", + " \n", + " true_vals_earth = truth_chain_plot[2*ii]\n", + " true_vals_jup = truth_chain_plot[2*ii+1]\n", " plot_truth = True\n", - "\n", - " fig, ax = plt.subplots(1, 2, figsize=(15, 3))\n", + " \n", + " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", - "\n", + " \n", " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", - "\n", - " nstep = np.arange(1, len(planet1_chain1) + 1, 1)\n", - "\n", + " \n", + " \n", + " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", + " \n", + " \n", " ax[0].plot(nstep, planet1_chain1)\n", " ax[0].plot(nstep, planet1_chain2)\n", - "\n", + " \n", " if plot_truth:\n", - " ax[0].axhline(y=true_vals_earth, color=\"r\", label=\"truth\")\n", - " ax[0].set_title(\"Jupiter\", fontsize=18)\n", - " ax[0].legend(fontsize=18)\n", - "\n", + " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", + " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", + " ax[0].legend(fontsize = 18)\n", + " \n", " ax[1].plot(nstep, planet2_chain1)\n", " ax[1].plot(nstep, planet2_chain2)\n", - "\n", + " \n", " if plot_truth:\n", - " ax[1].axhline(y=true_vals_jup, color=\"r\", label=\"truth\")\n", - " ax[1].set_title(\"Earth\", fontsize=18)\n", - " ax[1].legend(fontsize=18)\n", + " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", + " ax[1].set_title(\"Earth\", fontsize = 18)\n", + " ax[1].legend(fontsize = 18)\n", "\n", - " fig.suptitle(param, fontsize=18)\n", + " fig.suptitle(param, fontsize = 18)\n", " fig.tight_layout()\n", " plt.show()\n", "\n", - " return \"plotting complete\"\n", "\n", "\n", - "def a_from_Kepler3(period, M_tot):\n", - " period = period * 86400 # days to seconds\n", "\n", - " M_tot = M_tot * M_sun.value # solar masses to kg\n", + " return \"plotting complete\"\n", + "\n", "\n", - " a3 = (((constants.G.value) * M_tot) / (4 * np.pi**2)) * (period) ** 2.0\n", "\n", - " a = a3 ** (1 / 3)\n", "\n", - " a = a * 6.68459 * 10 ** (-12.0) # meters to AU\n", "\n", - " return a # in AUs" + "def a_from_Kepler3(period, M_tot):\n", + " period = period*86400 #days to seconds\n", + " \n", + "\n", + " M_tot = M_tot*M_sun.value #solar masses to kg\n", + " \n", + " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", + " \n", + " a = a3**(1/3)\n", + " \n", + " \n", + " \n", + " a = a * 6.68459*10**(-12.) # meters to AU\n", + " \n", + " return(a) #in AUs\n", + "\n", + "\n", + "\n" ] }, { @@ -1399,14 +1344,18 @@ "cell_type": "code", "execution_count": null, "id": "cfd9e0cd", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import arviz as az\n", "\n", - "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" + "az.summary(trace, var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"m_planet\",\n", + " \"Omega\",\n", + " \"omega\",\n", + " \"incl\"])\n" ] }, { @@ -1414,23 +1363,21 @@ "id": "6d8c580e", "metadata": {}, "source": [ - "#### Great, the r_hat statistitcs are all 1.0, suggesting that the chains are converged...\n", + "Great, the r_hat statistitcs are all 1.0, suggesting that the chains are converged...\n", "\n", - "#### ...we can also check this by eye by looking at the corner plots and the posterior chains!" + " ...we can also check this by eye by looking at the corner plots and the posterior chains!" ] }, { "cell_type": "code", "execution_count": null, "id": "6775ae19", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import corner\n", "\n", - "make_plots(trace, planet_params, 2)" + "make_plots(trace, planet_params, 2)\n" ] }, { @@ -1449,31 +1396,34 @@ "outputs": [], "source": [ "import pandas as pd\n", - "\n", - "# save simulated data as dataframe\n", + "#save simulated data as dataframe \n", "simulated_data_dic = {\n", " \"times_rv_observed\": x_rv,\n", " \"rv_observed\": y_rv,\n", " \"rv_err_observed\": y_rv_err,\n", + "\n", " \"times_rv_orbit\": times_rv,\n", " \"rv_orbit\": rv_orbit,\n", " \"rv_orbit_sum\": rv_orbit_sum,\n", + "\n", " \"times_astrometry_observed\": x_astrometry,\n", " \"ra_observed\": ra_data,\n", " \"dec_observed\": dec_data,\n", " \"ra_err_observed\": ra_err,\n", " \"dec_err_observed\": dec_err,\n", + "\n", " \"times_astrometry_orbit\": times_astrometry,\n", " \"ra_orbit\": ra_orbit,\n", " \"ra_orbit_sum\": ra_orbit_sum,\n", " \"dec_orbit\": dec_orbit,\n", - " \"dec_orbit_sum\": dec_orbit_sum,\n", - "}\n", + " \"dec_orbit_sum\": dec_orbit_sum\n", + " }\n", "\n", "\n", - "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient=\"index\")\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", "simulated_data = simulated_data.transpose()\n", - "simulated_data = simulated_data.fillna(value=np.nan)" + "simulated_data = simulated_data.fillna(value=np.nan)\n", + "\n" ] }, { @@ -1485,48 +1435,55 @@ "source": [ "def remove_nans(array):\n", " array = array[~pd.isnull(array)]\n", - "\n", + " \n", " return array\n", "\n", "\n", + "\n", + "\n", "def import_orbit_data(orbit_row):\n", " import math\n", - "\n", + " \n", " orbit_row = remove_nans(orbit_row)\n", " orbit_array = []\n", " for row in orbit_row:\n", " orbit_array.append([np.array(row[1]), row[0]])\n", "\n", " orbit_array = np.array(orbit_array)\n", - "\n", + " \n", " return orbit_array\n", "\n", "\n", - "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", - " import matplotlib.pyplot as plt\n", - "\n", - " # simulated RV observations\n", - " x_rv = simulated_data[\"times_rv_observed\"].values\n", - " y_rv = simulated_data[\"rv_observed\"].values\n", - " y_rv_err = simulated_data[\"rv_err_observed\"].values\n", - "\n", - " # simulated RV full orbits\n", - " x_rv_orbit = simulated_data[\"times_rv_orbit\"].values\n", - " y_rv_orbit = simulated_data[\"rv_orbit\"].values\n", "\n", - " # simulated astrometric observations\n", - " x_astrometry = simulated_data[\"times_astrometry_observed\"].values\n", - " y_ra = simulated_data[\"ra_observed\"].values\n", - " y_dec = simulated_data[\"dec_observed\"].values\n", - " y_ra_err = simulated_data[\"ra_err_observed\"].values\n", - " y_dec_err = simulated_data[\"dec_err_observed\"].values\n", "\n", - " # simulated astrometric full orbits\n", - " x_astrometry_orbit = simulated_data[\"times_astrometry_orbit\"].values\n", - " y_ra_orbit = simulated_data[\"ra_orbit\"].values\n", - " y_dec_orbit = simulated_data[\"dec_orbit\"].values\n", + "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", + " import matplotlib.pyplot as plt\n", "\n", - " # remove any nans from simulated data (created because of mismatching lengths)\n", + " \n", + " \n", + " #simulated RV observations\n", + " x_rv = simulated_data['times_rv_observed'].values\n", + " y_rv = simulated_data['rv_observed'].values\n", + " y_rv_err = simulated_data['rv_err_observed'].values\n", + " \n", + " #simulated RV full orbits\n", + " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", + " y_rv_orbit = simulated_data['rv_orbit'].values\n", + " \n", + " #simulated astrometric observations\n", + " x_astrometry = simulated_data['times_astrometry_observed'].values\n", + " y_ra = simulated_data['ra_observed'].values\n", + " y_dec = simulated_data['dec_observed'].values\n", + " y_ra_err = simulated_data['ra_err_observed'].values\n", + " y_dec_err = simulated_data['dec_err_observed'].values\n", + " \n", + " #simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", + " y_ra_orbit = simulated_data['ra_orbit'].values\n", + " y_dec_orbit = simulated_data['dec_orbit'].values\n", + " \n", + " \n", + " #remove any nans from simulated data (created because of mismatching lengths)\n", " x_rv = remove_nans(x_rv)\n", " y_rv = remove_nans(y_rv)\n", " y_rv_err = remove_nans(y_rv_err)\n", @@ -1540,34 +1497,43 @@ " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", - "\n", + " \n", + " \n", + " \n", + " \n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - " )\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", - "\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " \n", + " \n", + " \n", + " \n", + " \n", " # add NaN for plotting purpose on true model\n", - " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float(\"NaN\"))\n", - " y_ra_orbit = np.insert(y_ra_orbit, 10000, float(\"NaN\"), axis=0)\n", - " y_dec_orbit = np.insert(y_dec_orbit, 10000, float(\"NaN\"), axis=0)\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", + " \n", + " \n", + " \n", + "\n", "\n", + " \n", + " \n", " figsAs = []\n", " figsRV = []\n", " for n in range(0, nplanets):\n", + " \n", " ######\n", - " # make rv plots\n", - "\n", - " figRV, ax = plt.subplots(\n", - " 2, 1, figsize=[13, 6], gridspec_kw={\"height_ratios\": [3, 1]}\n", - " )\n", + " #make rv plots\n", + " \n", + " \n", + " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", " ax0, ax1 = ax[0], ax[1]\n", - "\n", + " \n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1578,68 +1544,86 @@ " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # Plot the data\n", - " ax0.errorbar(\n", - " x_rv,\n", - " y_rv - other,\n", - " yerr=y_rv_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - "\n", + " \n", " med_rv = []\n", " for an_x in x_rv:\n", " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", "\n", + "\n", + "\n", + " \n", " # plot the MCMC model\n", " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(t_rv, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax0.fill_between(\n", + " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", - "\n", + " \n", " # plot the true RA model\n", - " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color=\"k\", label=\"truth\")\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", - "\n", - " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " # plot residuals\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " #plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_rv,\n", - " y_rv - other - med_rv,\n", - " yerr=y_rv_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_rv, y_rv - other - med_rv,\n", + " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", " )\n", - "\n", + " \n", + " \n", " figRV.tight_layout()\n", " figRV.subplots_adjust(hspace=0.5)\n", - " figRV.suptitle(\"planet {0}\".format(n + 1), fontsize=33)\n", + " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", " figRV.tight_layout()\n", - "\n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", " figRV.show()\n", " ################\n", - " # plot astrometry\n", - "\n", - " figAs, ax = plt.subplots(\n", - " 2,\n", - " 2,\n", - " figsize=[18, 9],\n", - " sharey=\"row\",\n", - " gridspec_kw={\"height_ratios\": [3, 1]},\n", - " )\n", + " #plot astrometry\n", + " \n", + " \n", + "\n", + " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", "\n", " # Get the posterior median orbital parameters\n", @@ -1652,16 +1636,14 @@ " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # plot the DEC data\n", - " ax0.errorbar(\n", - " x_astrometry,\n", - " y_dec - other,\n", - " yerr=y_dec_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", "\n", " # plot the MCMC DEC model\n", " pred = np.percentile(\n", @@ -1669,78 +1651,55 @@ " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - "\n", + " \n", " med_dec = []\n", " for an_x in x_astrometry:\n", " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", + " \n", + "\n", "\n", " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax0.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", "\n", " # plot the true DEC model\n", - " ax0.plot(\n", - " x_astrometry_orbit, y_dec_orbit.T[n], color=\"k\", label=\"truth\"\n", - " )\n", + " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + " \n", + " \n", + " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", + " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", + " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + "\n", + "\n", "\n", - " ax0.arrow(\n", - " 0.5,\n", - " 0.0,\n", - " 0,\n", - " 1,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.0,\n", - " head_length=0.0,\n", - " color=\"k\",\n", - " ls=\"--\",\n", - " )\n", - " ax0.text(\n", - " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax0.transAxes\n", - " )\n", - " ax0.text(\n", - " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax0.transAxes\n", - " )\n", - " ax0.arrow(\n", - " 0.49,\n", - " 0.87,\n", - " -0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", - " ax0.arrow(\n", - " 0.51,\n", - " 0.87,\n", - " 0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize=18)\n", - "\n", - " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", + " \n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - "\n", - " # plot residuals\n", + " \n", + " \n", + " #plot residuals \n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_astrometry,\n", - " y_dec - other - med_dec,\n", - " yerr=y_dec_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_astrometry, y_dec - other - med_dec,\n", + " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", " )\n", "\n", + "\n", + "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1750,99 +1709,86 @@ " other = np.median(\n", " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", + " \n", + "\n", + "\n", + "\n", + "\n", "\n", " # Plot the RA data\n", - " ax2.errorbar(\n", - " x_astrometry,\n", - " y_ra - other,\n", - " yerr=y_ra_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", " [16, 50, 84],\n", - " axis=(0, 1),\n", + " axis=(0, 1), \n", " )\n", - "\n", + " \n", + " \n", " med_ra = []\n", " for an_x in x_astrometry:\n", " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", "\n", + "\n", " # plot the MCMC RA model\n", " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax2.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax2.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", - "\n", + " \n", " # plot the true RA model\n", - " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color=\"k\", label=\"truth\")\n", - "\n", - " ax2.arrow(\n", - " 0.5,\n", - " 0.0,\n", - " 0,\n", - " 1,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.0,\n", - " head_length=0.0,\n", - " color=\"k\",\n", - " ls=\"--\",\n", - " )\n", - " ax2.text(\n", - " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax2.transAxes\n", - " )\n", - " ax2.text(\n", - " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax2.transAxes\n", - " )\n", - " ax2.arrow(\n", - " 0.49,\n", - " 0.87,\n", - " -0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", - " ax2.arrow(\n", - " 0.51,\n", - " 0.87,\n", - " 0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + " \n", + " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", + " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", + " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", "\n", + "\n", + " \n", " ax2.legend(fontsize=9, loc=1)\n", " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize=18)\n", - "\n", - " ax3.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", + " \n", + " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " # plot residuals\n", + " \n", + " #plot residuals\n", " ax3.axhline(0.0, color=\"k\")\n", " ax3.errorbar(\n", - " x_astrometry,\n", - " y_ra - other - med_ra,\n", - " yerr=y_ra_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_astrometry, y_ra - other - med_ra,\n", + " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", " )\n", - "\n", + " \n", " figAs.tight_layout()\n", - " figAs.subplots_adjust(wspace=0.2)\n", - "\n", + " figAs.subplots_adjust(wspace=.2)\n", + " \n", + " \n", " figAs.show()\n", - "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", " return None" ] }, @@ -1850,22 +1796,18 @@ "cell_type": "code", "execution_count": null, "id": "b3941ebb", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ - "make_rv_astrometry_plots(simulated_data, trace, 2)" + "make_rv_astrometry_plots(simulated_data, trace, 2)\n" ] }, { "cell_type": "markdown", "id": "1f6c90cd", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ - "### And lastly, check how well our MCMC median +/- 1 sigma results agree with the simulated values" + "## And lastly, check how well our MCMC posterior median +/- 1 sigma results agree with the simulated values" ] }, { @@ -1876,7 +1818,9 @@ "outputs": [], "source": [ "def load_posterior_params(trace, nplanets):\n", - " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\n", + "\n", + " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", + "\n", "\n", " params_earth = defaultdict(list)\n", " params_jup = defaultdict(list)\n", @@ -1891,139 +1835,138 @@ " posterior_meds = defaultdict(list)\n", " posterior_errs = defaultdict(list)\n", " for param in parameters:\n", + "\n", + "\n", " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", "\n", - " planet_quantile = [\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.84),\n", - " ]\n", + " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", + "\n", + "\n", "\n", " planet_err = [\n", " planet_med - planet_quantile[0],\n", - " planet_quantile[1] - planet_med,\n", + " planet_quantile[1] - planet_med \n", " ]\n", "\n", - " # print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + "\n", + "\n", + " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", "\n", " posterior_meds[param] = np.round(planet_med, 3)\n", " posterior_errs[param] = np.round(planet_err, 3)\n", - " table_params[param].append(\n", - " str(np.round(planet_med, 3))\n", - " + \" (+\"\n", - " + str(np.round(planet_err[1], 3))\n", - " + \" -\"\n", - " + str(np.round(planet_err[0], 3))\n", - " + \")\"\n", - " )\n", + " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", + " str(np.round(planet_err[1], 3)) + \" -\" + \n", + " str(np.round(planet_err[0], 3)) + \")\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " table_params_all.append(table_params)\n", " posterior_meds_all.append(posterior_meds)\n", " posterior_errs_all.append(posterior_errs)\n", - " # print(\"\")\n", - " # print(\"\")\n", + " #print(\"\")\n", + " #print(\"\")\n", + "\n", + "\n", "\n", " return table_params_all, posterior_meds_all, posterior_errs_all\n", "\n", "\n", - "def make_comparison_figures(\n", - " input_params,\n", - " posterior_orbit_params,\n", - " keys_orbit_params,\n", - " posterior_meds,\n", - " posterior_errs,\n", - " nplanets,\n", - "):\n", + "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", + " posterior_meds, posterior_errs, nplanets):\n", + "\n", " import matplotlib\n", " import matplotlib.pyplot as plt\n", - "\n", + " \n", + " \n", " input_params = np.array(input_params)\n", - " input_params.T[-1] = input_params.T[-1] * 332946.07\n", + " input_params.T[-1] = input_params.T[-1]*m_sun\n", "\n", " comparisons = []\n", " for ii in range(0, nplanets):\n", - " fig, axs = plt.subplots(\n", - " 1, 2, figsize=[13, 6], gridspec_kw={\"width_ratios\": [1, 1]}\n", - " )\n", + " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", " plt.subplots_adjust(hspace=1)\n", - "\n", + " \n", " #### add comparison plots\n", - "\n", + " \n", + " \n", " input_mass = input_params[ii][-1]\n", " input_period = input_params[ii][0]\n", + " \n", + " xs = posterior_meds[ii]['m_planet']\n", + " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", "\n", - " xs = posterior_meds[ii][\"m_planet\"]\n", - " xs_err = np.array([posterior_errs[ii][\"m_planet\"]]).T\n", - "\n", - " ys = posterior_meds[ii][\"P\"]\n", - " ys_err = np.array([posterior_errs[ii][\"P\"]]).T\n", - "\n", - " color = \"#1c245f\"\n", - "\n", - " axs[0].errorbar(\n", - " xs,\n", - " ys,\n", - " xerr=xs_err,\n", - " yerr=ys_err,\n", - " marker=\"o\",\n", - " color=color,\n", - " markersize=13,\n", - " label=r\"med and 1$\\sigma$ errs\",\n", - " linewidth=3,\n", - " alpha=0.7,\n", - " )\n", + " ys = posterior_meds[ii]['P']\n", + " ys_err = np.array([posterior_errs[ii]['P']]).T\n", + "\n", + "\n", + "\n", + " color = '#1c245f'\n", + "\n", + " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", + " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", + " linewidth = 3, alpha = 0.7)\n", + "\n", + " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", "\n", - " axs[0].axvline(\n", - " x=input_mass, ymin=0, ymax=1, color=\"k\", ls=\"--\", label=\"simulated\"\n", - " )\n", - " axs[0].axhline(y=input_period, xmin=0, xmax=1, color=\"k\", ls=\"--\")\n", "\n", - " axs[0].set_xlabel(r\"mass [M$_{earth}$]\", fontsize=27)\n", - " axs[0].set_ylabel(r\"period [days]\", fontsize=27)\n", + " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", + " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", "\n", - " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", "\n", - " axs[0].legend(fontsize=13)\n", + " fig.suptitle('input and posterior parameters', fontsize = 27)\n", "\n", + " axs[0].legend(fontsize = 13)\n", + " \n", + " \n", + " \n", " #### add comparison tables\n", "\n", + "\n", " sim_params = input_params[ii]\n", " post_params = posterior_orbit_params[ii]\n", - " data = np.column_stack(\n", - " (\n", - " np.array(keys_orbit_params),\n", - " np.round(np.array(sim_params), 3),\n", - " np.array(list(post_params.values())),\n", - " )\n", - " )\n", + " data = np.column_stack((np.array(keys_orbit_params), \n", + " np.round(np.array(sim_params),3),\n", + " np.array(list(post_params.values()))))\n", "\n", - " labels = [\"param\", \"truth\", r\"MCMC med +/- 1$\\sigma$\"]\n", + " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", "\n", " df = pd.DataFrame(data, columns=labels)\n", "\n", - " table = axs[1].table(\n", - " cellText=df.values,\n", - " colLabels=df.columns,\n", - " colColours=[\"#D7E5F0\"] * 3,\n", - " loc=\"center\",\n", - " bbox=[0, 0, 1, 1],\n", - " colWidths=[0.2, 0.2, 0.6],\n", - " )\n", - " # for (row, col), cell in table.get_celld().items():\n", + " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", + " colColours =[\"#D7E5F0\"]*3, loc='center', \n", + " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", + " #for (row, col), cell in table.get_celld().items():\n", " # if (row == 0):\n", " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", "\n", + "\n", " table.auto_set_column_width(col=labels)\n", "\n", - " # ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", - " # t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", - " # t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", " axs[1].set_axis_off()\n", "\n", + "\n", " table.auto_set_font_size(False)\n", " table.set_fontsize(13)\n", "\n", - " comparisons.append(fig)\n", + " \n", + " \n", + " \n", + "\n", + "\n", "\n", + " comparisons.append(fig)\n", + " \n", + " \n", + " \n", + " \n", " return comparisons" ] }, @@ -2037,29 +1980,12 @@ "from collections import defaultdict\n", "\n", "\n", - "(\n", - " table_params_all,\n", - " posterior_meds_all,\n", - " posterior_errs_all,\n", - ") = load_posterior_params(trace, 2)\n", - "\n", - "planet_param_keys = [\n", - " \"period\",\n", - " \"ecc\",\n", - " r\"T$_\\mathrm{per}$\",\n", - " r\"$\\omega$\",\n", - " r\"$\\Omega$\",\n", - " \"inc\",\n", - " \"mass\",\n", - "]\n", - "make_comparison_figures(\n", - " [planet_params[1], planet_params[0]],\n", - " table_params_all,\n", - " planet_param_keys,\n", - " posterior_meds_all,\n", - " posterior_errs_all,\n", - " 2,\n", - ")" + "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", + "load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", + "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", + " posterior_meds_all, posterior_errs_all, 2)" ] }, { From 8557c88a3517df49801a28e802006b2869a032bb Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Mar 2023 17:28:10 +0000 Subject: [PATCH 4/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/astrometry_rv_case_study.ipynb | 1740 +++++++++-------- 1 file changed, 907 insertions(+), 833 deletions(-) diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb index af760fdd..6a29e62f 100644 --- a/docs/tutorials/astrometry_rv_case_study.ipynb +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -12,7 +12,9 @@ "cell_type": "code", "execution_count": null, "id": "4abd3961", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -20,8 +22,7 @@ "\n", "\n", "xo.utils.docs_setup()\n", - "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", - "\n" + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")" ] }, { @@ -29,7 +30,7 @@ "id": "9a87ff47", "metadata": {}, "source": [ - "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023.\n", "\n", "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", "\n", @@ -49,12 +50,14 @@ { "cell_type": "markdown", "id": "93ceb2a9", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## (Simulated) Data\n", "\n", "\n", - "First, we need to simulate our dataset. \n", + "First, we need to simulate our dataset.\n", "\n", "\n", "### Let's define 2 helpful functions in order to simulate our data...." @@ -68,27 +71,24 @@ "outputs": [], "source": [ "def create_orbit(n_planets, orbit_params):\n", - " \n", - " '''\n", + " \"\"\"\n", " inputs:\n", " - n_planets = number of planets\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", - " '''\n", - " \n", + " - m_planet = mass of planet in Earth masses\n", + " \"\"\"\n", "\n", - " \n", " period = []\n", " ecc = []\n", - " Tper = [] \n", - " omega = [] \n", - " Omega = [] \n", + " Tper = []\n", + " omega = []\n", + " Omega = []\n", " incl = []\n", " m_planet = []\n", " for ii in range(0, n_planets):\n", @@ -101,15 +101,15 @@ " m_planet.append(orbit_params[ii][6])\n", "\n", " orbit = xo.orbits.KeplerianOrbit(\n", - " period = period,\n", - " ecc = ecc, \n", - " t_periastron = Tper, \n", - " omega = omega, \n", - " Omega = Omega, \n", - " incl = incl,\n", - " m_planet = m_planet) \n", - " \n", - " \n", + " period=period,\n", + " ecc=ecc,\n", + " t_periastron=Tper,\n", + " omega=omega,\n", + " Omega=Omega,\n", + " incl=incl,\n", + " m_planet=m_planet,\n", + " )\n", + "\n", " return orbit" ] }, @@ -120,24 +120,24 @@ "metadata": {}, "outputs": [], "source": [ - "np.random.seed(979) \n", + "np.random.seed(979)\n", + "\n", "\n", "def simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra,\n", " sigma_dec,\n", " plx,\n", " orbit_params,\n", - " times_observed_rv = None,\n", - " times_observed_astrometry = None,\n", - " t_dur_rv = None, \n", - " n_obs_rv = None,\n", - " t_dur_astrometry = None, \n", - " n_obs_astrometry = None \n", - " ):\n", - "\n", - " '''\n", + " times_observed_rv=None,\n", + " times_observed_astrometry=None,\n", + " t_dur_rv=None,\n", + " n_obs_rv=None,\n", + " t_dur_astrometry=None,\n", + " n_obs_astrometry=None,\n", + "):\n", + " \"\"\"\n", " inputs:\n", " - orbit = keplerian orbit generated by exoplanet code\n", " - n_planets = number of planets\n", @@ -147,230 +147,231 @@ " - plx = parallax of system\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", + " - m_planet = mass of planet in Earth masses\n", " - times_observed_rv = (default=None) the observed times for RV observations\n", " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", - " \n", - " \n", - " '''\n", - " \n", + "\n", + "\n", + " \"\"\"\n", + "\n", " if n_planets == 1:\n", " orbit_params = [orbit_params]\n", - " \n", - " \n", - " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", - " if times_observed_rv != None:\n", - " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", - "\n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined\n", + "\n", + " # if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(\n", + " np.min(times_observed_rv), np.max(times_observed_rv), 10000\n", + " )\n", + "\n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_rv = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_rv:\n", " start_time_rv = orbit_params[ii][2]\n", "\n", + " # define RV observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv + t_dur_rv, 10000)\n", + " times_observed_rv = np.linspace(\n", + " start_time_rv, start_time_rv + t_dur_rv, n_obs_rv\n", + " )\n", "\n", - " \n", - " #define RV observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", - " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", - " \n", - "\n", + " # if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(\n", + " np.min(times_observed_astrometry),\n", + " np.max(times_observed_astrometry),\n", + " 10000,\n", + " )\n", "\n", - " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", - " if times_observed_astrometry != None:\n", - " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", - " \n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined \n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_astrometry = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_astrometry:\n", - " start_time_astrometry = orbit_params[ii][2] \n", - "\n", - "\n", - " #define astrometry observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", - " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", - "\n", + " start_time_astrometry = orbit_params[ii][2]\n", + "\n", + " # define astrometry observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " 10000,\n", + " )\n", + " times_observed_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " n_obs_astrometry,\n", + " )\n", "\n", - " \n", " orbit = create_orbit(n_planets, orbit_params)\n", - " \n", + "\n", " rv = orbit.get_radial_velocity(times_rv)\n", " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", - " \n", + "\n", " rv_orbit = rv.eval()\n", "\n", " if n_planets > 1:\n", - " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", + " rv_orbit_sum = np.sum(rv_orbit, axis=1)\n", " else:\n", " rv_orbit_sum = rv_orbit\n", - " \n", "\n", " rv_observed = rv_observed.eval()\n", " if n_planets > 1:\n", - " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", + " rv_observed_sum = np.sum(rv_observed, axis=1)\n", " else:\n", " rv_observed_sum = rv_observed\n", "\n", - "\n", - " #-----------\n", - " #-----------\n", + " # -----------\n", + " # -----------\n", "\n", " # determine and print the star position at desired times\n", " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", "\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", - " #pos = tt.sum(pos, axis=-1)\n", - "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " rho = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", + "\n", " rho, theta = rho.eval(), theta.eval()\n", - " \n", - " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", + " # rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", " rho_orbit = rho\n", - " \n", + "\n", " if n_planets > 1:\n", - " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", + " rho_orbit_sum = np.sum(rho_orbit, axis=1)\n", " else:\n", " rho_orbit_sum = rho_orbit\n", "\n", " theta_orbit = theta\n", - " \n", + "\n", " if n_planets > 1:\n", - " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", + " theta_orbit_sum = np.sum(theta_orbit, axis=1)\n", " else:\n", " theta_orbit_sum = theta_orbit\n", - " \n", - " \n", + "\n", " # when summing over theta, position angle, we have to careful because position\n", - " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract\n", " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", " # in the correct range. Be careful though if modeling more than 2 planets, this\n", " # doesn't completely solve the problem!\n", - " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", - " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", - "\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2 * np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2 * np.pi\n", "\n", " # determine and print the star position at desired times\n", - " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", - "\n", - "\n", - " #pos = tt.sum(pos, axis=-1)\n", + " pos_observed = theano.function(\n", + " [], orbit.get_star_position(times_observed_astrometry, plx)\n", + " )()\n", "\n", - " x_obs,y_obs,z_obs = pos_observed\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", + " x_obs, y_obs, z_obs = pos_observed\n", "\n", " # calculate rho and theta\n", - " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", - " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", - " \n", - " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs**2 + y_obs**2)) # arcsec\n", + " theta_observed = tt.squeeze(\n", + " tt.arctan2(y_obs, x_obs)\n", + " ) # radians between [-pi, pi]\n", "\n", + " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", "\n", - " \n", " if n_planets > 1:\n", - " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", + " rho_observed_sum = np.sum(rho_observed, axis=1)\n", " else:\n", " rho_observed_sum = rho_observed\n", "\n", - " \n", " if n_planets > 1:\n", - " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", + " theta_observed_sum = np.sum(theta_observed, axis=1)\n", " else:\n", " theta_observed_sum = theta_observed\n", - " \n", - " \n", + "\n", " #######\n", " #######\n", " #######\n", "\n", " # convert rho and theta into ra and dec\n", - " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", "\n", " if n_planets > 1:\n", - " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", + " ra_orbit_sum = np.sum(ra_orbit, axis=1)\n", " else:\n", " ra_orbit_sum = ra_orbit\n", "\n", - " \n", " if n_planets > 1:\n", - " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", + " dec_orbit_sum = np.sum(dec_orbit, axis=1)\n", " else:\n", " dec_orbit_sum = dec_orbit\n", "\n", - " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", "\n", - "\n", - "\n", - "\n", - " \n", - " #----------- \n", - " #----------- \n", - " #----------- simulate data\n", - " #----------- \n", - " #----------- \n", - "\n", + " # -----------\n", + " # -----------\n", + " # ----------- simulate data\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", " rv_sim = rv_observed\n", - " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", + " rv_sim_sum = np.sum(rv_observed, axis=1) + np.random.normal(\n", + " 0, sigma_rv, len(rv_observed)\n", + " )\n", " else:\n", " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", " rv_sim_sum = rv_sim\n", "\n", - " #-----------\n", - " #-----------\n", - "\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", - " ra_sim = ra_observed \n", - " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", - " \n", + " ra_sim = ra_observed\n", + " ra_sim_sum = np.sum(ra_sim, axis=1) + np.random.normal(\n", + " 0, sigma_ra, len(ra_observed)\n", + " )\n", + "\n", " else:\n", " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", " ra_sim_sum = ra_sim\n", - " \n", - " \n", - " \n", + "\n", " if n_planets > 1:\n", - " dec_sim = dec_observed \n", - " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", - " \n", + " dec_sim = dec_observed\n", + " dec_sim_sum = np.sum(dec_sim, axis=1) + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", + "\n", " else:\n", - " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " dec_sim = dec_observed + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", " dec_sim_sum = dec_sim\n", "\n", - " \n", - " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", + " times = [\n", + " times_rv,\n", + " times_observed_rv,\n", + " times_astrometry,\n", + " times_observed_astrometry,\n", + " ]\n", " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", - " \n", - " \n", + "\n", " return [times, rv_results, ra_results, dec_results]" ] }, @@ -398,10 +399,20 @@ "source": [ "import numpy as np\n", "\n", - "parallax = 0.1 #as \n", - "n_planets=2\n", - "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", - " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" + "parallax = 0.1 # as\n", + "n_planets = 2\n", + "planet_params = [\n", + " [365.0, 0.0167, 100.0, 1.7959438003021653, 0.0, np.pi / 4, 3e-06],\n", + " [\n", + " 4327.0,\n", + " 0.0484,\n", + " 500.0,\n", + " -1.4957471689591395,\n", + " 1.7523105690023069,\n", + " np.pi / 4 + 0.022759,\n", + " 0.000955,\n", + " ],\n", + "]" ] }, { @@ -422,28 +433,28 @@ "cell_type": "code", "execution_count": null, "id": "55938572", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "from aesara_theano_fallback import aesara as theano\n", "import theano.tensor as tt\n", "\n", "\n", - "#add gaia observing times\n", + "# add gaia observing times\n", "times_observed_astrometry_gaia = []\n", "t_0 = 0\n", - "for ii in range(t_0, t_0+3600):\n", - " if ii % (3600/100) == 0:\n", + "for ii in range(t_0, t_0 + 3600):\n", + " if ii % (3600 / 100) == 0:\n", " times_observed_astrometry_gaia.append(ii)\n", "\n", "\n", - "\n", - "#add THE observing times\n", + "# add THE observing times\n", "times_observed_rv = []\n", - "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "t_0 = 3300 # Gaia first light (2014) ~9 years before THE first expected light (2023)\n", "add_data = True\n", - "for ii in range(t_0, t_0+3600):\n", - "\n", + "for ii in range(t_0, t_0 + 3600):\n", " if ii % 180 == 0:\n", " if add_data:\n", " add_data = False\n", @@ -454,72 +465,66 @@ " times_observed_rv.append(ii)\n", "\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "sigma_rv = 0.3\n", "\n", - "sigma_ra_gaia = 6e-5 #60 micro-as\n", - "sigma_dec_gaia = 6e-5 #60 micro-as\n", - "\n", + "sigma_ra_gaia = 6e-5 # 60 micro-as\n", + "sigma_dec_gaia = 6e-5 # 60 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_gaia,\n", " sigma_dec_gaia,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_gaia\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_gaia,\n", + ")\n", "\n", "\n", - "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", - "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", - "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", - "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", + "[\n", + " [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + " [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + " [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + " [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum],\n", + "] = (times, rv_results, ra_results, dec_results)\n", "\n", "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", "\n", "\n", - "\n", - "\n", - "\n", "# we assume that the first roman observation\n", "# will occur ~5 years after the last Gaia observation\n", - "t_1 = times_observed_astrometry_gaia[-1]+1800\n", + "t_1 = times_observed_astrometry_gaia[-1] + 1800\n", "\n", "# add roman observing times\n", "times_observed_astrometry_roman = []\n", - "for ii in range(t_1, t_1+(10*365)):\n", + "for ii in range(t_1, t_1 + (10 * 365)):\n", " if ii % 90 == 0:\n", - " times_observed_astrometry_roman.append(ii) \n", - "\n", + " times_observed_astrometry_roman.append(ii)\n", "\n", - "sigma_ra_roman = 5e-6 #5 micro-as\n", - "sigma_dec_roman = 5e-6 #5 micro-as\n", "\n", + "sigma_ra_roman = 5e-6 # 5 micro-as\n", + "sigma_dec_roman = 5e-6 # 5 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_roman,\n", " sigma_dec_roman,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_roman\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_roman,\n", + ")\n", "\n", "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", "\n", - "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", + "times_observed_astrometry = np.append(\n", + " times_observed_astrometry, times[3], axis=0\n", + ")\n", "\n", "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", @@ -532,7 +537,7 @@ "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", "\n", "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", - "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)" ] }, { @@ -552,36 +557,40 @@ "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib\n", - "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", - "matplotlib.rc('text', usetex=True)\n", + "\n", + "matplotlib.rc(\"font\", **{\"family\": \"serif\", \"serif\": [\"Computer Modern\"]})\n", + "matplotlib.rc(\"text\", usetex=True)\n", + "\n", "\n", "def plot_rv_signal(\n", - " n_planets, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv):\n", - " \n", - " \n", - " \n", - " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", - " fig.suptitle(\"RV Signal\", fontsize = 45)\n", + " n_planets,\n", + " rv_orbit,\n", + " rv_orbit_sum,\n", + " rv_sim,\n", + " rv_sim_sum,\n", + " times_rv,\n", + " times_observed_rv,\n", + "):\n", + " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", + " fig.suptitle(\"RV Signal\", fontsize=45)\n", " ax1 = ax\n", - " \n", - " \n", - " \n", "\n", - " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", + " ax1.plot(\n", + " times_observed_rv,\n", + " rv_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " alpha=0.3,\n", + " )\n", "\n", - " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", - " ax1.legend(fontsize = 27, loc = 2)\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize=27)\n", + " ax1.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -593,53 +602,55 @@ "outputs": [], "source": [ "def plot_astrometry_signal(\n", - " n_planets, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " n_planets,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry):\n", - " \n", - " \n", - "\n", - " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", - " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", + " times_observed_astrometry,\n", + "):\n", + " fig, ax = plt.subplots(2, 1, figsize=[13, 10], sharex=\"col\", sharey=\"row\")\n", + " fig.suptitle(\"Astrometric Signal\", fontsize=45)\n", " ax2 = ax[0]\n", " ax3 = ax[1]\n", - " \n", - "\n", "\n", - " \n", - " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.plot(\n", + " times_observed_astrometry,\n", + " ra_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", " for tick in ax3.get_xticklabels():\n", " tick.set_rotation(30)\n", - " ax2.legend(fontsize = 27, loc = 2)\n", - " \n", - " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", - " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax3.legend(fontsize = 27, loc = 2)\n", - "\n", - "\n", + " ax2.legend(fontsize=27, loc=2)\n", + "\n", + " ax3.plot(\n", + " times_observed_astrometry,\n", + " dec_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax3.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", "\n", - "\n", - "\n", " #######\n", " #######\n", " #######\n", "\n", - "\n", - " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", + " fig, ax = plt.subplots(1, 1, figsize=[10, 10])\n", "\n", " # plot simulated data\n", - " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", + " ax.plot(ra_sim_sum, dec_sim_sum, \"o\", color=\"k\", label=\"combined signal\")\n", "\n", " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", @@ -649,7 +660,7 @@ " ax.set_title(\"initial orbit\")\n", " ax.legend()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -665,46 +676,46 @@ "cell_type": "code", "execution_count": null, "id": "82e18e64", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_rv_signal(\n", - " 2, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv)\n", - " \n", - " \n" + " 2, rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum, times_rv, times_observed_rv\n", + ")" ] }, { "cell_type": "code", "execution_count": null, "id": "623b9c2d", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_astrometry_signal(\n", - " 2, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " 2,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry)" + " times_observed_astrometry,\n", + ")" ] }, { "cell_type": "markdown", "id": "99110cc0", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Ok, now we have a simulated dataset -- let's model it!\n", "\n", @@ -725,9 +736,7 @@ " print(\"------------\")\n", "\n", " with pm.Model() as model:\n", - " \n", - " \n", - " #log normal prior on period around estimates\n", + " # log normal prior on period around estimates\n", " logP = pm.Uniform(\n", " \"logP\",\n", " lower=0,\n", @@ -735,37 +744,33 @@ " shape=2,\n", " testval=np.log(periods),\n", " )\n", - " \n", - "\n", "\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", "\n", - "\n", " ## wide uniform prior on t_periastron\n", " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", - " \n", - " \n", + "\n", " # Wide normal prior for semi-amplitude\n", - " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", - " \n", + " logK = pm.Uniform(\n", + " \"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks)\n", + " )\n", + "\n", " K = pm.Deterministic(\"K\", tt.exp(logK))\n", - " \n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", "\n", - "\n", - "\n", - " \n", " # Jitter & a quadratic RV trend\n", - " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + " # logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", "\n", " # Then we define the orbit\n", - " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", + " orbit = xo.orbits.KeplerianOrbit(\n", + " period=P, t_periastron=tperi, ecc=ecc, omega=omega\n", + " )\n", "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", @@ -776,8 +781,6 @@ " # Sum over planets and add the background to get the full model\n", " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", - " \n", - "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", "\n", @@ -786,19 +789,17 @@ "\n", " # Finally add in the observation model. This next line adds a new contribution\n", " # to the log probability of the PyMC3 model\n", - " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " # err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", " err = y_rv_err\n", " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", "\n", - "\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", " map_soln = pmx.optimize(start=map_soln)\n", "\n", - "\n", - " #return the max a-posteriori solution\n", + " # return the max a-posteriori solution\n", " return map_soln" ] }, @@ -806,16 +807,18 @@ "cell_type": "code", "execution_count": null, "id": "4d6faa40", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import pymc3 as pm\n", "import pymc3_ext as pmx\n", "\n", - "np.random.seed(979) \n", + "np.random.seed(979)\n", "\n", "\n", - "#rename variables in more consistent way for modeling\n", + "# rename variables in more consistent way for modeling\n", "x_rv = np.array(times_observed_rv)\n", "y_rv = rv_sim_sum\n", "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", @@ -829,25 +832,29 @@ "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", "\n", "\n", - " \n", - " \n", "# make a fine grid that spans the observation window for plotting purposes\n", - "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + "t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + ")\n", "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", "# for predicted orbits\n", - "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + ")\n", "\n", "\n", - "#minimize on RV data\n", + "# minimize on RV data\n", "periods_real = [4327.0, 365.0]\n", "periods_guess = []\n", "\n", - "#let's assume that THE can recover the period within 10% of the true value\n", + "# let's assume that THE can recover the period within 10% of the true value\n", "for per in periods_real:\n", - " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", + " periods_guess.append(\n", + " per + np.random.uniform(low=-0.1 * per, high=0.1 * per)\n", + " )\n", "\n", - "print('periods used as starting point = ', periods_guess)\n", + "print(\"periods used as starting point = \", periods_guess)\n", "periods_guess = np.array(periods_guess)\n", "\n", "\n", @@ -859,7 +866,9 @@ { "cell_type": "markdown", "id": "efc4dcc2", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Now let's define out joint astrometry and radial velocity model:" ] @@ -871,25 +880,40 @@ "metadata": {}, "outputs": [], "source": [ - "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", - " m_sun = 333030 #earth masses\n", - " \n", - " P_RV = np.array(rv_map_soln['P'])\n", - " K_RV = np.array(rv_map_soln['K'])\n", - " tperi_RV = np.array(rv_map_soln['tperi'])\n", - " ecc_RV = np.array(rv_map_soln['ecc'])\n", - " omega_RV = np.array(rv_map_soln['omega'])\n", - " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", + "def minimize_both(\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", + "):\n", + " m_sun = 333030 # earth masses\n", + "\n", + " P_RV = np.array(rv_map_soln[\"P\"])\n", + " K_RV = np.array(rv_map_soln[\"K\"])\n", + " tperi_RV = np.array(rv_map_soln[\"tperi\"])\n", + " ecc_RV = np.array(rv_map_soln[\"ecc\"])\n", + " omega_RV = np.array(rv_map_soln[\"omega\"])\n", + " min_masses_RV = (\n", + " xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value * 317.83\n", + " ) # in m_earth\n", " phase_RV = determine_phase(P_RV, tperi_RV)\n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - "\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", " print(\"RV Solutions\")\n", " print(\"------------\")\n", @@ -899,98 +923,106 @@ " print(\"eccentricity: \", ecc_RV)\n", " print(\"omega: \", omega_RV)\n", "\n", - " print('')\n", + " print(\"\")\n", " print(\"minimizing joint model solutions pre-MCMC\")\n", " print(\"------------\")\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", - " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", + " inc_test_vals = np.array(np.radians([5.0, 25.0, 45.0, 65.0, 85.0]))\n", " model, map_soln, logp = [], [], []\n", " for inc in inc_test_vals:\n", - " mass_test_vals = min_masses_RV/np.sin(inc)\n", - " print('')\n", + " mass_test_vals = min_masses_RV / np.sin(inc)\n", + " print(\"\")\n", " print(\"trying inclination = \" + str(np.degrees(inc)))\n", " print(\"mass test val = \" + str(mass_test_vals))\n", " print(\"------------\")\n", - " \n", - " \n", - "\n", - "\n", "\n", " def get_model():\n", " with pm.Model() as model:\n", - "\n", - "\n", " # Below we will run a version of this model where a measurement of parallax is provided\n", " # The measurement is in milliarcsec\n", " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", - " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", + " \"m_plx\", mu=parallax * 1000, sd=10, testval=parallax * 1000\n", " )\n", " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", "\n", - "\n", " # We expect the period to be around that found from just the RVs,\n", " # so we'll set a broad prior on logP\n", - " \n", + "\n", " logP = pm.Uniform(\n", - " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", + " \"logP\",\n", + " lower=0,\n", + " upper=np.log(2 * P_RV),\n", + " testval=np.log(P_RV),\n", + " shape=2,\n", " )\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", - " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", - " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", - " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecs = pmx.UnitDisk(\n", + " \"ecs\",\n", + " shape=(2, 2),\n", + " testval=np.array(\n", + " [\n", + " np.sqrt(ecc_RV) * np.cos(omega_RV),\n", + " np.sqrt(ecc_RV) * np.sin(omega_RV),\n", + " ]\n", + " ),\n", + " )\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", - " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", - "\n", - " \n", - " \n", + " xo.eccentricity.kipping13(\n", + " \"ecc_prior\", fixed=True, observed=ecc\n", + " )\n", "\n", " # Omegas are co-dependent, so sample them with variables Omega_plus\n", - " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and\n", " # Omega_minus is (Omega_0 - Omega_1)/2\n", - " \n", + "\n", " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", - " \n", - " \n", - " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", - " (Omega_plus - Omega_minus)] )\n", - " \n", "\n", - " Omega = pm.Deterministic(\"Omega\", Omega) \n", - " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", - " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", - " \n", + " Omega = tt.concatenate(\n", + " [(Omega_plus + Omega_minus), (Omega_plus - Omega_minus)]\n", + " )\n", "\n", + " Omega = pm.Deterministic(\"Omega\", Omega)\n", + " Omega_sum = pm.Deterministic(\n", + " \"Omega_sum\", ((Omega_plus) * 2) % np.pi\n", + " )\n", + " Omega_diff = pm.Deterministic(\n", + " \"Omega_diff\", ((Omega_minus) * 2) % np.pi\n", + " )\n", "\n", - " \n", " # For these orbits, it can also be better to fit for a phase angle\n", " # (relative to a reference time) instead of the time of periasteron\n", " # passage directly\n", " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", - " \n", "\n", - " \n", " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", - " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", + " log_m = pm.Uniform(\n", + " \"log_m\",\n", + " lower=-1,\n", + " upper=np.log(100 * mass_test_vals),\n", + " testval=np.log(mass_test_vals),\n", + " shape=2,\n", + " )\n", " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", - " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", - "\n", + " m_planet_fit = pm.Deterministic(\n", + " \"m_planet_fit\", m_planet / m_sun\n", + " )\n", "\n", - " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", + " cos_incl = pm.Uniform(\n", + " \"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2\n", + " )\n", " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", - " \n", - "\n", "\n", - " \n", - " \n", " # Set up the orbit\n", " orbit = xo.orbits.KeplerianOrbit(\n", " t_periastron=tperi,\n", @@ -999,13 +1031,10 @@ " ecc=ecc,\n", " omega=omega,\n", " Omega=Omega,\n", - " m_planet = m_planet_fit,\n", - " plx=plx\n", + " m_planet=m_planet_fit,\n", + " plx=plx,\n", " )\n", "\n", - "\n", - " \n", - " \n", " # Add a function for computing the full astrometry model\n", " def get_astrometry_model(t, name=\"\"):\n", " # First the astrometry induced by the planets\n", @@ -1013,54 +1042,56 @@ " # determine and print the star position at desired times\n", " pos = orbit.get_star_position(t, plx)\n", "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", - " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", - " \n", - " #rhos, thetas = get_star_relative_angles(t, plx)\n", - " \n", - " \n", - " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", - " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", - " \n", + " rhos = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", + " thetas = tt.squeeze(\n", + " tt.arctan2(y, x)\n", + " ) # radians between [-pi, pi]\n", + "\n", + " # rhos, thetas = get_star_relative_angles(t, plx)\n", + "\n", + " dec = pm.Deterministic(\n", + " \"dec\" + name, rhos * np.cos(thetas)\n", + " ) # X is north\n", + " ra = pm.Deterministic(\n", + " \"ra\" + name, rhos * np.sin(thetas)\n", + " ) # Y is east\n", + "\n", " # Sum over planets to get the full model\n", - " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", - " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", - " \n", + " dec_model = pm.Deterministic(\n", + " \"dec_model\" + name, tt.sum(dec, axis=-1)\n", + " )\n", + " ra_model = pm.Deterministic(\n", + " \"ra_model\" + name, tt.sum(ra, axis=-1)\n", + " )\n", "\n", - " \n", " return dec_model, ra_model\n", "\n", - " \n", " # Define the astrometry model at the observed times\n", " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", "\n", " # Also define the model on a fine grid as computed above (for plotting)\n", - " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", - "\n", - " \n", - "\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(\n", + " t_fine, name=\"_fine\"\n", + " )\n", "\n", - " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", - " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " # dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " # ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", " dec_tot_err = dec_err\n", " ra_tot_err = ra_err\n", "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec\n", + " pm.Normal(\n", + " \"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data\n", + " )\n", + " pm.Normal(\n", + " \"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data\n", + " )\n", "\n", - " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", - " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", - " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", - "\n", - "\n", - " \n", - " \n", " # ADD RV MODEL\n", "\n", - "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", " # First the RVs induced by the planets\n", @@ -1068,7 +1099,9 @@ " pm.Deterministic(\"vrad\" + name, vrad)\n", "\n", " # Sum over planets to get the full model\n", - " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", + " return pm.Deterministic(\n", + " \"rv_model\" + name, tt.sum(vrad, axis=-1)\n", + " )\n", "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", @@ -1084,24 +1117,20 @@ " # Optimize to find the initial parameters\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", - " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + " # map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", "\n", " map_soln = pmx.optimize(map_soln)\n", "\n", - "\n", " return model, map_soln\n", "\n", - "\n", - "\n", " a_model, a_map_soln = get_model()\n", - " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", - " print('log likelihood = ' + str(a_logp))\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis=0)\n", + " print(\"log likelihood = \" + str(a_logp))\n", "\n", " model.append(a_model)\n", " map_soln.append(a_map_soln)\n", " logp.append(a_logp)\n", "\n", - " \n", " best_index = 0\n", " for index in range(0, len(model)):\n", " if logp[index] >= logp[best_index]:\n", @@ -1111,31 +1140,26 @@ " the_map_soln = map_soln[best_index]\n", " the_logp = logp[best_index]\n", "\n", - "\n", - "\n", " return the_model, the_map_soln, the_logp\n", "\n", "\n", - "\n", "def determine_phase(P, t_periastron):\n", " phase = (2 * np.pi * t_periastron) / P\n", " return phase\n", "\n", "\n", - "\n", "def model_both(model, map_soln, tune_steps, draw_steps):\n", - " print('Joint RV + Astometry Minimization Solutions:')\n", + " print(\"Joint RV + Astometry Minimization Solutions:\")\n", " print(\"------------\")\n", "\n", - " print('m_planet: ' + str(map_soln['m_planet']))\n", - " print('P: ' + str(map_soln['P']))\n", - " print('incl: ' + str(map_soln['incl']))\n", - " print('Omega: ' + str(map_soln['Omega']))\n", - " print('tperi: ' + str(map_soln['tperi']))\n", - " print('ecc: ' + str(map_soln['ecc']))\n", - " print('omega: ' + str(map_soln['omega']))\n", - " print('plx: ' + str(map_soln['plx']))\n", - "\n", + " print(\"m_planet: \" + str(map_soln[\"m_planet\"]))\n", + " print(\"P: \" + str(map_soln[\"P\"]))\n", + " print(\"incl: \" + str(map_soln[\"incl\"]))\n", + " print(\"Omega: \" + str(map_soln[\"Omega\"]))\n", + " print(\"tperi: \" + str(map_soln[\"tperi\"]))\n", + " print(\"ecc: \" + str(map_soln[\"ecc\"]))\n", + " print(\"omega: \" + str(map_soln[\"omega\"]))\n", + " print(\"plx: \" + str(map_soln[\"plx\"]))\n", "\n", " with model:\n", " trace = pmx.sample(\n", @@ -1163,18 +1187,24 @@ "\n", "################\n", "################\n", - "#minimize on joint model\n", + "# minimize on joint model\n", "joint_model, joint_map_soln, joint_logp = minimize_both(\n", - " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", - " ra_data, ra_err, dec_data, dec_err, parallax\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", ")\n", "\n", "\n", - "\n", - "\n", "################\n", "################\n", - "#run full MCMC\n", + "# run full MCMC\n", "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" ] }, @@ -1193,143 +1223,168 @@ "metadata": {}, "outputs": [], "source": [ - "matplotlib.rc('text', usetex=False)\n", + "matplotlib.rc(\"text\", usetex=False)\n", + "\n", "\n", "def make_plots(trace, orbit_params, n_planets):\n", - " m_sun = 333030 #earth masses\n", + " m_sun = 333030 # earth masses\n", "\n", " [orbit_params_earth, orbit_params_jup] = orbit_params\n", "\n", - "\n", - " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", - " inclination_earth, m_earth] = orbit_params_earth\n", - "\n", - "\n", - " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", - " inclination_jup, m_jup] = orbit_params_jup\n", - "\n", - "\n", - " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", - " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", - " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", - " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", - " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", - "\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", - " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", - " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " [\n", + " P_earth,\n", + " e_earth,\n", + " Tper_earth,\n", + " omega_earth,\n", + " Omega_earth,\n", + " inclination_earth,\n", + " m_earth,\n", + " ] = orbit_params_earth\n", + "\n", + " [\n", + " P_jup,\n", + " e_jup,\n", + " Tper_jup,\n", + " omega_jup,\n", + " Omega_jup,\n", + " inclination_jup,\n", + " m_jup,\n", + " ] = orbit_params_jup\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0 + m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0 + m_jup)\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " Omega_earth - Omega_jup,\n", + " Omega_earth + Omega_jup,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " 0.1,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " ]\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " ]\n", "\n", " # plot the table summarizing MCMC results\n", " az.summary(\n", " trace,\n", - " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", - " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", + " var_names=[\n", + " \"P\",\n", + " \"tperi\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_plus\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"ecc\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " ],\n", " )\n", "\n", - "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", " # plot the corner plots\n", " _ = corner.corner(\n", - " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", - " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", - " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", - " truths = truths, truth_color = \"#03003a\"\n", + " trace,\n", + " var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " \"tperi\",\n", + " ],\n", + " quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True,\n", + " title_kwargs={\"fontsize\": 13},\n", + " truths=truths,\n", + " truth_color=\"#03003a\",\n", " )\n", "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the chains\n", " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", " for ii in range(0, len(parameters)):\n", " plot_truth = False\n", " param = parameters[ii]\n", - " \n", - " true_vals_earth = truth_chain_plot[2*ii]\n", - " true_vals_jup = truth_chain_plot[2*ii+1]\n", + "\n", + " true_vals_earth = truth_chain_plot[2 * ii]\n", + " true_vals_jup = truth_chain_plot[2 * ii + 1]\n", " plot_truth = True\n", - " \n", - " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", + "\n", + " fig, ax = plt.subplots(1, 2, figsize=(15, 3))\n", " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", - " \n", + "\n", " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", - " \n", - " \n", - " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", - " \n", - " \n", + "\n", + " nstep = np.arange(1, len(planet1_chain1) + 1, 1)\n", + "\n", " ax[0].plot(nstep, planet1_chain1)\n", " ax[0].plot(nstep, planet1_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", - " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", - " ax[0].legend(fontsize = 18)\n", - " \n", + " ax[0].axhline(y=true_vals_earth, color=\"r\", label=\"truth\")\n", + " ax[0].set_title(\"Jupiter\", fontsize=18)\n", + " ax[0].legend(fontsize=18)\n", + "\n", " ax[1].plot(nstep, planet2_chain1)\n", " ax[1].plot(nstep, planet2_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", - " ax[1].set_title(\"Earth\", fontsize = 18)\n", - " ax[1].legend(fontsize = 18)\n", + " ax[1].axhline(y=true_vals_jup, color=\"r\", label=\"truth\")\n", + " ax[1].set_title(\"Earth\", fontsize=18)\n", + " ax[1].legend(fontsize=18)\n", "\n", - " fig.suptitle(param, fontsize = 18)\n", + " fig.suptitle(param, fontsize=18)\n", " fig.tight_layout()\n", " plt.show()\n", "\n", + " return \"plotting complete\"\n", "\n", "\n", + "def a_from_Kepler3(period, M_tot):\n", + " period = period * 86400 # days to seconds\n", "\n", - " return \"plotting complete\"\n", - "\n", + " M_tot = M_tot * M_sun.value # solar masses to kg\n", "\n", + " a3 = (((constants.G.value) * M_tot) / (4 * np.pi**2)) * (period) ** 2.0\n", "\n", + " a = a3 ** (1 / 3)\n", "\n", + " a = a * 6.68459 * 10 ** (-12.0) # meters to AU\n", "\n", - "def a_from_Kepler3(period, M_tot):\n", - " period = period*86400 #days to seconds\n", - " \n", - "\n", - " M_tot = M_tot*M_sun.value #solar masses to kg\n", - " \n", - " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", - " \n", - " a = a3**(1/3)\n", - " \n", - " \n", - " \n", - " a = a * 6.68459*10**(-12.) # meters to AU\n", - " \n", - " return(a) #in AUs\n", - "\n", - "\n", - "\n" + " return a # in AUs" ] }, { @@ -1344,18 +1399,14 @@ "cell_type": "code", "execution_count": null, "id": "cfd9e0cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import arviz as az\n", "\n", - "az.summary(trace, var_names=[\n", - " \"P\",\n", - " \"ecc\",\n", - " \"m_planet\",\n", - " \"Omega\",\n", - " \"omega\",\n", - " \"incl\"])\n" + "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" ] }, { @@ -1372,12 +1423,14 @@ "cell_type": "code", "execution_count": null, "id": "6775ae19", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import corner\n", "\n", - "make_plots(trace, planet_params, 2)\n" + "make_plots(trace, planet_params, 2)" ] }, { @@ -1396,34 +1449,31 @@ "outputs": [], "source": [ "import pandas as pd\n", - "#save simulated data as dataframe \n", + "\n", + "# save simulated data as dataframe\n", "simulated_data_dic = {\n", " \"times_rv_observed\": x_rv,\n", " \"rv_observed\": y_rv,\n", " \"rv_err_observed\": y_rv_err,\n", - "\n", " \"times_rv_orbit\": times_rv,\n", " \"rv_orbit\": rv_orbit,\n", " \"rv_orbit_sum\": rv_orbit_sum,\n", - "\n", " \"times_astrometry_observed\": x_astrometry,\n", " \"ra_observed\": ra_data,\n", " \"dec_observed\": dec_data,\n", " \"ra_err_observed\": ra_err,\n", " \"dec_err_observed\": dec_err,\n", - "\n", " \"times_astrometry_orbit\": times_astrometry,\n", " \"ra_orbit\": ra_orbit,\n", " \"ra_orbit_sum\": ra_orbit_sum,\n", " \"dec_orbit\": dec_orbit,\n", - " \"dec_orbit_sum\": dec_orbit_sum\n", - " }\n", + " \"dec_orbit_sum\": dec_orbit_sum,\n", + "}\n", "\n", "\n", - "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient=\"index\")\n", "simulated_data = simulated_data.transpose()\n", - "simulated_data = simulated_data.fillna(value=np.nan)\n", - "\n" + "simulated_data = simulated_data.fillna(value=np.nan)" ] }, { @@ -1435,55 +1485,48 @@ "source": [ "def remove_nans(array):\n", " array = array[~pd.isnull(array)]\n", - " \n", - " return array\n", - "\n", "\n", + " return array\n", "\n", "\n", "def import_orbit_data(orbit_row):\n", " import math\n", - " \n", + "\n", " orbit_row = remove_nans(orbit_row)\n", " orbit_array = []\n", " for row in orbit_row:\n", " orbit_array.append([np.array(row[1]), row[0]])\n", "\n", " orbit_array = np.array(orbit_array)\n", - " \n", - " return orbit_array\n", - "\n", "\n", + " return orbit_array\n", "\n", "\n", "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", " import matplotlib.pyplot as plt\n", "\n", - " \n", - " \n", - " #simulated RV observations\n", - " x_rv = simulated_data['times_rv_observed'].values\n", - " y_rv = simulated_data['rv_observed'].values\n", - " y_rv_err = simulated_data['rv_err_observed'].values\n", - " \n", - " #simulated RV full orbits\n", - " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", - " y_rv_orbit = simulated_data['rv_orbit'].values\n", - " \n", - " #simulated astrometric observations\n", - " x_astrometry = simulated_data['times_astrometry_observed'].values\n", - " y_ra = simulated_data['ra_observed'].values\n", - " y_dec = simulated_data['dec_observed'].values\n", - " y_ra_err = simulated_data['ra_err_observed'].values\n", - " y_dec_err = simulated_data['dec_err_observed'].values\n", - " \n", - " #simulated astrometric full orbits\n", - " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", - " y_ra_orbit = simulated_data['ra_orbit'].values\n", - " y_dec_orbit = simulated_data['dec_orbit'].values\n", - " \n", - " \n", - " #remove any nans from simulated data (created because of mismatching lengths)\n", + " # simulated RV observations\n", + " x_rv = simulated_data[\"times_rv_observed\"].values\n", + " y_rv = simulated_data[\"rv_observed\"].values\n", + " y_rv_err = simulated_data[\"rv_err_observed\"].values\n", + "\n", + " # simulated RV full orbits\n", + " x_rv_orbit = simulated_data[\"times_rv_orbit\"].values\n", + " y_rv_orbit = simulated_data[\"rv_orbit\"].values\n", + "\n", + " # simulated astrometric observations\n", + " x_astrometry = simulated_data[\"times_astrometry_observed\"].values\n", + " y_ra = simulated_data[\"ra_observed\"].values\n", + " y_dec = simulated_data[\"dec_observed\"].values\n", + " y_ra_err = simulated_data[\"ra_err_observed\"].values\n", + " y_dec_err = simulated_data[\"dec_err_observed\"].values\n", + "\n", + " # simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data[\"times_astrometry_orbit\"].values\n", + " y_ra_orbit = simulated_data[\"ra_orbit\"].values\n", + " y_dec_orbit = simulated_data[\"dec_orbit\"].values\n", + "\n", + " # remove any nans from simulated data (created because of mismatching lengths)\n", " x_rv = remove_nans(x_rv)\n", " y_rv = remove_nans(y_rv)\n", " y_rv_err = remove_nans(y_rv_err)\n", @@ -1497,43 +1540,34 @@ " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " # add NaN for plotting purpose on true model\n", - " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", - " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", - " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", - " \n", - " \n", - " \n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", + " # add NaN for plotting purpose on true model\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float(\"NaN\"))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float(\"NaN\"), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float(\"NaN\"), axis=0)\n", "\n", - " \n", - " \n", " figsAs = []\n", " figsRV = []\n", " for n in range(0, nplanets):\n", - " \n", " ######\n", - " #make rv plots\n", - " \n", - " \n", - " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", + " # make rv plots\n", + "\n", + " figRV, ax = plt.subplots(\n", + " 2, 1, figsize=[13, 6], gridspec_kw={\"height_ratios\": [3, 1]}\n", + " )\n", " ax0, ax1 = ax[0], ax[1]\n", - " \n", + "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1544,86 +1578,68 @@ " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # Plot the data\n", - " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax0.errorbar(\n", + " x_rv,\n", + " y_rv - other,\n", + " yerr=y_rv_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_rv = []\n", " for an_x in x_rv:\n", " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", "\n", - "\n", - "\n", - " \n", " # plot the MCMC model\n", " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_rv, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "\n", + " # plot the true RA model\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color=\"k\", label=\"truth\")\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", - " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - "\n", "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " \n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_rv, y_rv - other - med_rv,\n", - " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", + " x_rv,\n", + " y_rv - other - med_rv,\n", + " yerr=y_rv_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", - " \n", + "\n", " figRV.tight_layout()\n", " figRV.subplots_adjust(hspace=0.5)\n", - " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", + " figRV.suptitle(\"planet {0}\".format(n + 1), fontsize=33)\n", " figRV.tight_layout()\n", - " \n", - " \n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", + "\n", " figRV.show()\n", " ################\n", - " #plot astrometry\n", - " \n", - " \n", - "\n", - " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", + " # plot astrometry\n", + "\n", + " figAs, ax = plt.subplots(\n", + " 2,\n", + " 2,\n", + " figsize=[18, 9],\n", + " sharey=\"row\",\n", + " gridspec_kw={\"height_ratios\": [3, 1]},\n", + " )\n", " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", "\n", " # Get the posterior median orbital parameters\n", @@ -1636,14 +1652,16 @@ " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the DEC data\n", - " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + " ax0.errorbar(\n", + " x_astrometry,\n", + " y_dec - other,\n", + " yerr=y_dec_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " # plot the MCMC DEC model\n", " pred = np.percentile(\n", @@ -1651,55 +1669,78 @@ " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_dec = []\n", " for an_x in x_astrometry:\n", " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", - " \n", - "\n", "\n", " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", "\n", " # plot the true DEC model\n", - " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - " \n", - " \n", - " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", - " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", - " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - "\n", - "\n", + " ax0.plot(\n", + " x_astrometry_orbit, y_dec_orbit.T[n], color=\"k\", label=\"truth\"\n", + " )\n", "\n", + " ax0.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax0.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax0.transAxes\n", + " )\n", + " ax0.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax0.transAxes\n", + " )\n", + " ax0.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax0.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", - " \n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - " \n", - " \n", - " #plot residuals \n", + "\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_astrometry, y_dec - other - med_dec,\n", - " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_dec - other - med_dec,\n", + " yerr=y_dec_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", "\n", - "\n", - "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1709,86 +1750,99 @@ " other = np.median(\n", " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", - " \n", - "\n", - "\n", - "\n", - "\n", "\n", " # Plot the RA data\n", - " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax2.errorbar(\n", + " x_astrometry,\n", + " y_ra - other,\n", + " yerr=y_ra_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", " [16, 50, 84],\n", - " axis=(0, 1), \n", + " axis=(0, 1),\n", " )\n", - " \n", - " \n", + "\n", " med_ra = []\n", " for an_x in x_astrometry:\n", " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", "\n", - "\n", " # plot the MCMC RA model\n", " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax2.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax2.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - " \n", - " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", - " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", - " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", "\n", + " # plot the true RA model\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color=\"k\", label=\"truth\")\n", + "\n", + " ax2.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax2.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax2.transAxes\n", + " )\n", + " ax2.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax2.transAxes\n", + " )\n", + " ax2.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax2.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", - " \n", " ax2.legend(fontsize=9, loc=1)\n", " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", - " \n", - " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax3.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax3.axhline(0.0, color=\"k\")\n", " ax3.errorbar(\n", - " x_astrometry, y_ra - other - med_ra,\n", - " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_ra - other - med_ra,\n", + " yerr=y_ra_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", + "\n", " figAs.tight_layout()\n", - " figAs.subplots_adjust(wspace=.2)\n", - " \n", - " \n", + " figAs.subplots_adjust(wspace=0.2)\n", + "\n", " figAs.show()\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "\n", - " \n", + "\n", " return None" ] }, @@ -1796,16 +1850,20 @@ "cell_type": "code", "execution_count": null, "id": "b3941ebb", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ - "make_rv_astrometry_plots(simulated_data, trace, 2)\n" + "make_rv_astrometry_plots(simulated_data, trace, 2)" ] }, { "cell_type": "markdown", "id": "1f6c90cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## And lastly, check how well our MCMC posterior median +/- 1 sigma results agree with the simulated values" ] @@ -1818,9 +1876,7 @@ "outputs": [], "source": [ "def load_posterior_params(trace, nplanets):\n", - "\n", - " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", - "\n", + " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\n", "\n", " params_earth = defaultdict(list)\n", " params_jup = defaultdict(list)\n", @@ -1835,138 +1891,139 @@ " posterior_meds = defaultdict(list)\n", " posterior_errs = defaultdict(list)\n", " for param in parameters:\n", - "\n", - "\n", " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", "\n", - " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", - "\n", - "\n", + " planet_quantile = [\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84),\n", + " ]\n", "\n", " planet_err = [\n", " planet_med - planet_quantile[0],\n", - " planet_quantile[1] - planet_med \n", + " planet_quantile[1] - planet_med,\n", " ]\n", "\n", - "\n", - "\n", - " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + " # print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", "\n", " posterior_meds[param] = np.round(planet_med, 3)\n", " posterior_errs[param] = np.round(planet_err, 3)\n", - " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", - " str(np.round(planet_err[1], 3)) + \" -\" + \n", - " str(np.round(planet_err[0], 3)) + \")\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " table_params[param].append(\n", + " str(np.round(planet_med, 3))\n", + " + \" (+\"\n", + " + str(np.round(planet_err[1], 3))\n", + " + \" -\"\n", + " + str(np.round(planet_err[0], 3))\n", + " + \")\"\n", + " )\n", "\n", " table_params_all.append(table_params)\n", " posterior_meds_all.append(posterior_meds)\n", " posterior_errs_all.append(posterior_errs)\n", - " #print(\"\")\n", - " #print(\"\")\n", - "\n", - "\n", + " # print(\"\")\n", + " # print(\"\")\n", "\n", " return table_params_all, posterior_meds_all, posterior_errs_all\n", "\n", "\n", - "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", - " posterior_meds, posterior_errs, nplanets):\n", - "\n", + "def make_comparison_figures(\n", + " input_params,\n", + " posterior_orbit_params,\n", + " keys_orbit_params,\n", + " posterior_meds,\n", + " posterior_errs,\n", + " nplanets,\n", + "):\n", " import matplotlib\n", " import matplotlib.pyplot as plt\n", - " \n", - " \n", + "\n", " input_params = np.array(input_params)\n", - " input_params.T[-1] = input_params.T[-1]*m_sun\n", + " input_params.T[-1] = input_params.T[-1] * m_sun\n", "\n", " comparisons = []\n", " for ii in range(0, nplanets):\n", - " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", + " fig, axs = plt.subplots(\n", + " 1, 2, figsize=[13, 6], gridspec_kw={\"width_ratios\": [1, 1]}\n", + " )\n", " plt.subplots_adjust(hspace=1)\n", - " \n", + "\n", " #### add comparison plots\n", - " \n", - " \n", + "\n", " input_mass = input_params[ii][-1]\n", " input_period = input_params[ii][0]\n", - " \n", - " xs = posterior_meds[ii]['m_planet']\n", - " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", - "\n", - " ys = posterior_meds[ii]['P']\n", - " ys_err = np.array([posterior_errs[ii]['P']]).T\n", - "\n", "\n", + " xs = posterior_meds[ii][\"m_planet\"]\n", + " xs_err = np.array([posterior_errs[ii][\"m_planet\"]]).T\n", + "\n", + " ys = posterior_meds[ii][\"P\"]\n", + " ys_err = np.array([posterior_errs[ii][\"P\"]]).T\n", + "\n", + " color = \"#1c245f\"\n", + "\n", + " axs[0].errorbar(\n", + " xs,\n", + " ys,\n", + " xerr=xs_err,\n", + " yerr=ys_err,\n", + " marker=\"o\",\n", + " color=color,\n", + " markersize=13,\n", + " label=r\"med and 1$\\sigma$ errs\",\n", + " linewidth=3,\n", + " alpha=0.7,\n", + " )\n", "\n", - " color = '#1c245f'\n", - "\n", - " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", - " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", - " linewidth = 3, alpha = 0.7)\n", - "\n", - " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", - " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", - "\n", + " axs[0].axvline(\n", + " x=input_mass, ymin=0, ymax=1, color=\"k\", ls=\"--\", label=\"simulated\"\n", + " )\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color=\"k\", ls=\"--\")\n", "\n", - " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", - " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", + " axs[0].set_xlabel(r\"mass [M$_{earth}$]\", fontsize=27)\n", + " axs[0].set_ylabel(r\"period [days]\", fontsize=27)\n", "\n", + " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", "\n", - " fig.suptitle('input and posterior parameters', fontsize = 27)\n", + " axs[0].legend(fontsize=13)\n", "\n", - " axs[0].legend(fontsize = 13)\n", - " \n", - " \n", - " \n", " #### add comparison tables\n", "\n", - "\n", " sim_params = input_params[ii]\n", " post_params = posterior_orbit_params[ii]\n", - " data = np.column_stack((np.array(keys_orbit_params), \n", - " np.round(np.array(sim_params),3),\n", - " np.array(list(post_params.values()))))\n", + " data = np.column_stack(\n", + " (\n", + " np.array(keys_orbit_params),\n", + " np.round(np.array(sim_params), 3),\n", + " np.array(list(post_params.values())),\n", + " )\n", + " )\n", "\n", - " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", + " labels = [\"param\", \"truth\", r\"MCMC med +/- 1$\\sigma$\"]\n", "\n", " df = pd.DataFrame(data, columns=labels)\n", "\n", - " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", - " colColours =[\"#D7E5F0\"]*3, loc='center', \n", - " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", - " #for (row, col), cell in table.get_celld().items():\n", + " table = axs[1].table(\n", + " cellText=df.values,\n", + " colLabels=df.columns,\n", + " colColours=[\"#D7E5F0\"] * 3,\n", + " loc=\"center\",\n", + " bbox=[0, 0, 1, 1],\n", + " colWidths=[0.2, 0.2, 0.6],\n", + " )\n", + " # for (row, col), cell in table.get_celld().items():\n", " # if (row == 0):\n", " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", "\n", - "\n", " table.auto_set_column_width(col=labels)\n", "\n", - " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", - " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", - " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " # ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " # t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " # t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", " axs[1].set_axis_off()\n", "\n", - "\n", " table.auto_set_font_size(False)\n", " table.set_fontsize(13)\n", "\n", - " \n", - " \n", - " \n", - "\n", - "\n", - "\n", " comparisons.append(fig)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " return comparisons" ] }, @@ -1980,12 +2037,29 @@ "from collections import defaultdict\n", "\n", "\n", - "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", - "load_posterior_params(trace, 2)\n", - "\n", - "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", - "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", - " posterior_meds_all, posterior_errs_all, 2)" + "(\n", + " table_params_all,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + ") = load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = [\n", + " \"period\",\n", + " \"ecc\",\n", + " r\"T$_\\mathrm{per}$\",\n", + " r\"$\\omega$\",\n", + " r\"$\\Omega$\",\n", + " \"inc\",\n", + " \"mass\",\n", + "]\n", + "make_comparison_figures(\n", + " [planet_params[1], planet_params[0]],\n", + " table_params_all,\n", + " planet_param_keys,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + " 2,\n", + ")" ] }, { From 099699cce417ae4baf4a8e72151a273c14172cc9 Mon Sep 17 00:00:00 2001 From: dyahalomi Date: Wed, 15 Mar 2023 11:36:10 -0600 Subject: [PATCH 5/6] clear the notebook output with nbstripout --- docs/tutorials/astrometry_rv_case_study.ipynb | 1740 ++++++++--------- 1 file changed, 833 insertions(+), 907 deletions(-) diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb index 6a29e62f..af760fdd 100644 --- a/docs/tutorials/astrometry_rv_case_study.ipynb +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -12,9 +12,7 @@ "cell_type": "code", "execution_count": null, "id": "4abd3961", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -22,7 +20,8 @@ "\n", "\n", "xo.utils.docs_setup()\n", - "print(f\"exoplanet.__version__ = '{xo.__version__}'\")" + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", + "\n" ] }, { @@ -30,7 +29,7 @@ "id": "9a87ff47", "metadata": {}, "source": [ - "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023.\n", + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", "\n", "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", "\n", @@ -50,14 +49,12 @@ { "cell_type": "markdown", "id": "93ceb2a9", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## (Simulated) Data\n", "\n", "\n", - "First, we need to simulate our dataset.\n", + "First, we need to simulate our dataset. \n", "\n", "\n", "### Let's define 2 helpful functions in order to simulate our data...." @@ -71,24 +68,27 @@ "outputs": [], "source": [ "def create_orbit(n_planets, orbit_params):\n", - " \"\"\"\n", + " \n", + " '''\n", " inputs:\n", " - n_planets = number of planets\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", " - period = period in days\n", - " - ecc = eccentricity\n", + " - ecc = eccentricity \n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node\n", + " - Omega = longitude of ascending node \n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses\n", - " \"\"\"\n", + " - m_planet = mass of planet in Earth masses \n", + " '''\n", + " \n", "\n", + " \n", " period = []\n", " ecc = []\n", - " Tper = []\n", - " omega = []\n", - " Omega = []\n", + " Tper = [] \n", + " omega = [] \n", + " Omega = [] \n", " incl = []\n", " m_planet = []\n", " for ii in range(0, n_planets):\n", @@ -101,15 +101,15 @@ " m_planet.append(orbit_params[ii][6])\n", "\n", " orbit = xo.orbits.KeplerianOrbit(\n", - " period=period,\n", - " ecc=ecc,\n", - " t_periastron=Tper,\n", - " omega=omega,\n", - " Omega=Omega,\n", - " incl=incl,\n", - " m_planet=m_planet,\n", - " )\n", - "\n", + " period = period,\n", + " ecc = ecc, \n", + " t_periastron = Tper, \n", + " omega = omega, \n", + " Omega = Omega, \n", + " incl = incl,\n", + " m_planet = m_planet) \n", + " \n", + " \n", " return orbit" ] }, @@ -120,24 +120,24 @@ "metadata": {}, "outputs": [], "source": [ - "np.random.seed(979)\n", - "\n", + "np.random.seed(979) \n", "\n", "def simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra,\n", " sigma_dec,\n", " plx,\n", " orbit_params,\n", - " times_observed_rv=None,\n", - " times_observed_astrometry=None,\n", - " t_dur_rv=None,\n", - " n_obs_rv=None,\n", - " t_dur_astrometry=None,\n", - " n_obs_astrometry=None,\n", - "):\n", - " \"\"\"\n", + " times_observed_rv = None,\n", + " times_observed_astrometry = None,\n", + " t_dur_rv = None, \n", + " n_obs_rv = None,\n", + " t_dur_astrometry = None, \n", + " n_obs_astrometry = None \n", + " ):\n", + "\n", + " '''\n", " inputs:\n", " - orbit = keplerian orbit generated by exoplanet code\n", " - n_planets = number of planets\n", @@ -147,231 +147,230 @@ " - plx = parallax of system\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", " - period = period in days\n", - " - ecc = eccentricity\n", + " - ecc = eccentricity \n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node\n", + " - Omega = longitude of ascending node \n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses\n", + " - m_planet = mass of planet in Earth masses \n", " - times_observed_rv = (default=None) the observed times for RV observations\n", " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", - "\n", - "\n", - " \"\"\"\n", - "\n", + " \n", + " \n", + " '''\n", + " \n", " if n_planets == 1:\n", " orbit_params = [orbit_params]\n", - "\n", - " # if times_observed_rv != None then just define times_rv and linspace between min and max\n", - " if times_observed_rv != None:\n", - " times_rv = np.linspace(\n", - " np.min(times_observed_rv), np.max(times_observed_rv), 10000\n", - " )\n", - "\n", - " # define start_time_RV, which will be used to determine the times to simulate\n", - " # if times_observed_RV and times_observed_astrometry aren't defined\n", + " \n", + " \n", + " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", + "\n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_rv = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_rv:\n", " start_time_rv = orbit_params[ii][2]\n", "\n", - " # define RV observation time\n", - " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_rv = np.linspace(start_time_rv, start_time_rv + t_dur_rv, 10000)\n", - " times_observed_rv = np.linspace(\n", - " start_time_rv, start_time_rv + t_dur_rv, n_obs_rv\n", - " )\n", "\n", - " # if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", - " if times_observed_astrometry != None:\n", - " times_astrometry = np.linspace(\n", - " np.min(times_observed_astrometry),\n", - " np.max(times_observed_astrometry),\n", - " 10000,\n", - " )\n", + " \n", + " #define RV observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", + " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", + " \n", + "\n", "\n", - " # define start_time_RV, which will be used to determine the times to simulate\n", - " # if times_observed_RV and times_observed_astrometry aren't defined\n", + " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", + " \n", + " #define start_time_RV, which will be used to determine the times to simulate \n", + " #if times_observed_RV and times_observed_astrometry aren't defined \n", " else:\n", " start_time_astrometry = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_astrometry:\n", - " start_time_astrometry = orbit_params[ii][2]\n", - "\n", - " # define astrometry observation time\n", - " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_astrometry = np.linspace(\n", - " start_time_astrometry,\n", - " start_time_astrometry + t_dur_astrometry,\n", - " 10000,\n", - " )\n", - " times_observed_astrometry = np.linspace(\n", - " start_time_astrometry,\n", - " start_time_astrometry + t_dur_astrometry,\n", - " n_obs_astrometry,\n", - " )\n", + " start_time_astrometry = orbit_params[ii][2] \n", "\n", - " orbit = create_orbit(n_planets, orbit_params)\n", "\n", + " #define astrometry observation time\n", + " #100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", + " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", + "\n", + "\n", + " \n", + " orbit = create_orbit(n_planets, orbit_params)\n", + " \n", " rv = orbit.get_radial_velocity(times_rv)\n", " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", - "\n", + " \n", " rv_orbit = rv.eval()\n", "\n", " if n_planets > 1:\n", - " rv_orbit_sum = np.sum(rv_orbit, axis=1)\n", + " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", " else:\n", " rv_orbit_sum = rv_orbit\n", + " \n", "\n", " rv_observed = rv_observed.eval()\n", " if n_planets > 1:\n", - " rv_observed_sum = np.sum(rv_observed, axis=1)\n", + " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", " else:\n", " rv_observed_sum = rv_observed\n", "\n", - " # -----------\n", - " # -----------\n", + "\n", + " #-----------\n", + " #-----------\n", "\n", " # determine and print the star position at desired times\n", " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", "\n", - " # pos = tt.sum(pos, axis=-1)\n", "\n", - " x, y, z = pos\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x,y,z = pos\n", + "\n", "\n", " # calculate rho and theta\n", - " rho = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", + " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - "\n", + " \n", " rho, theta = rho.eval(), theta.eval()\n", + " \n", + " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", - " # rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", " rho_orbit = rho\n", - "\n", + " \n", " if n_planets > 1:\n", - " rho_orbit_sum = np.sum(rho_orbit, axis=1)\n", + " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", " else:\n", " rho_orbit_sum = rho_orbit\n", "\n", " theta_orbit = theta\n", - "\n", + " \n", " if n_planets > 1:\n", - " theta_orbit_sum = np.sum(theta_orbit, axis=1)\n", + " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", " else:\n", " theta_orbit_sum = theta_orbit\n", - "\n", + " \n", + " \n", " # when summing over theta, position angle, we have to careful because position\n", - " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract\n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", " # in the correct range. Be careful though if modeling more than 2 planets, this\n", " # doesn't completely solve the problem!\n", - " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2 * np.pi\n", - " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2 * np.pi\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", + "\n", "\n", " # determine and print the star position at desired times\n", - " pos_observed = theano.function(\n", - " [], orbit.get_star_position(times_observed_astrometry, plx)\n", - " )()\n", + " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", "\n", - " # pos = tt.sum(pos, axis=-1)\n", "\n", - " x_obs, y_obs, z_obs = pos_observed\n", + " #pos = tt.sum(pos, axis=-1)\n", + "\n", + " x_obs,y_obs,z_obs = pos_observed\n", "\n", - " # calculate rho and theta\n", - " rho_observed = tt.squeeze(tt.sqrt(x_obs**2 + y_obs**2)) # arcsec\n", - " theta_observed = tt.squeeze(\n", - " tt.arctan2(y_obs, x_obs)\n", - " ) # radians between [-pi, pi]\n", "\n", + " # calculate rho and theta\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", + " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", + " \n", " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", "\n", + "\n", + " \n", " if n_planets > 1:\n", - " rho_observed_sum = np.sum(rho_observed, axis=1)\n", + " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", " else:\n", " rho_observed_sum = rho_observed\n", "\n", + " \n", " if n_planets > 1:\n", - " theta_observed_sum = np.sum(theta_observed, axis=1)\n", + " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", " else:\n", " theta_observed_sum = theta_observed\n", - "\n", + " \n", + " \n", " #######\n", " #######\n", " #######\n", "\n", " # convert rho and theta into ra and dec\n", - " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", "\n", " if n_planets > 1:\n", - " ra_orbit_sum = np.sum(ra_orbit, axis=1)\n", + " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", " else:\n", " ra_orbit_sum = ra_orbit\n", "\n", + " \n", " if n_planets > 1:\n", - " dec_orbit_sum = np.sum(dec_orbit, axis=1)\n", + " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", " else:\n", " dec_orbit_sum = dec_orbit\n", "\n", - " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", "\n", - " # -----------\n", - " # -----------\n", - " # ----------- simulate data\n", - " # -----------\n", - " # -----------\n", + "\n", + "\n", + "\n", + " \n", + " #----------- \n", + " #----------- \n", + " #----------- simulate data\n", + " #----------- \n", + " #----------- \n", + "\n", "\n", " if n_planets > 1:\n", " rv_sim = rv_observed\n", - " rv_sim_sum = np.sum(rv_observed, axis=1) + np.random.normal(\n", - " 0, sigma_rv, len(rv_observed)\n", - " )\n", + " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", " else:\n", " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", " rv_sim_sum = rv_sim\n", "\n", - " # -----------\n", - " # -----------\n", + " #-----------\n", + " #-----------\n", "\n", - " if n_planets > 1:\n", - " ra_sim = ra_observed\n", - " ra_sim_sum = np.sum(ra_sim, axis=1) + np.random.normal(\n", - " 0, sigma_ra, len(ra_observed)\n", - " )\n", "\n", + " if n_planets > 1:\n", + " ra_sim = ra_observed \n", + " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", + " \n", " else:\n", " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", " ra_sim_sum = ra_sim\n", - "\n", + " \n", + " \n", + " \n", " if n_planets > 1:\n", - " dec_sim = dec_observed\n", - " dec_sim_sum = np.sum(dec_sim, axis=1) + np.random.normal(\n", - " 0, sigma_dec, len(dec_observed)\n", - " )\n", - "\n", + " dec_sim = dec_observed \n", + " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " \n", " else:\n", - " dec_sim = dec_observed + np.random.normal(\n", - " 0, sigma_dec, len(dec_observed)\n", - " )\n", + " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", " dec_sim_sum = dec_sim\n", "\n", - " times = [\n", - " times_rv,\n", - " times_observed_rv,\n", - " times_astrometry,\n", - " times_observed_astrometry,\n", - " ]\n", + " \n", + " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", - "\n", + " \n", + " \n", " return [times, rv_results, ra_results, dec_results]" ] }, @@ -399,20 +398,10 @@ "source": [ "import numpy as np\n", "\n", - "parallax = 0.1 # as\n", - "n_planets = 2\n", - "planet_params = [\n", - " [365.0, 0.0167, 100.0, 1.7959438003021653, 0.0, np.pi / 4, 3e-06],\n", - " [\n", - " 4327.0,\n", - " 0.0484,\n", - " 500.0,\n", - " -1.4957471689591395,\n", - " 1.7523105690023069,\n", - " np.pi / 4 + 0.022759,\n", - " 0.000955,\n", - " ],\n", - "]" + "parallax = 0.1 #as \n", + "n_planets=2\n", + "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", + " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" ] }, { @@ -433,28 +422,28 @@ "cell_type": "code", "execution_count": null, "id": "55938572", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "from aesara_theano_fallback import aesara as theano\n", "import theano.tensor as tt\n", "\n", "\n", - "# add gaia observing times\n", + "#add gaia observing times\n", "times_observed_astrometry_gaia = []\n", "t_0 = 0\n", - "for ii in range(t_0, t_0 + 3600):\n", - " if ii % (3600 / 100) == 0:\n", + "for ii in range(t_0, t_0+3600):\n", + " if ii % (3600/100) == 0:\n", " times_observed_astrometry_gaia.append(ii)\n", "\n", "\n", - "# add THE observing times\n", + "\n", + "#add THE observing times\n", "times_observed_rv = []\n", - "t_0 = 3300 # Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", "add_data = True\n", - "for ii in range(t_0, t_0 + 3600):\n", + "for ii in range(t_0, t_0+3600):\n", + "\n", " if ii % 180 == 0:\n", " if add_data:\n", " add_data = False\n", @@ -465,66 +454,72 @@ " times_observed_rv.append(ii)\n", "\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "sigma_rv = 0.3\n", "\n", - "sigma_ra_gaia = 6e-5 # 60 micro-as\n", - "sigma_dec_gaia = 6e-5 # 60 micro-as\n", + "sigma_ra_gaia = 6e-5 #60 micro-as\n", + "sigma_dec_gaia = 6e-5 #60 micro-as\n", + "\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra_gaia,\n", " sigma_dec_gaia,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv=times_observed_rv,\n", - " times_observed_astrometry=times_observed_astrometry_gaia,\n", - ")\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_gaia\n", + " )\n", "\n", "\n", - "[\n", - " [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", - " [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", - " [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", - " [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum],\n", - "] = (times, rv_results, ra_results, dec_results)\n", + "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", "\n", "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", "\n", "\n", + "\n", + "\n", + "\n", "# we assume that the first roman observation\n", "# will occur ~5 years after the last Gaia observation\n", - "t_1 = times_observed_astrometry_gaia[-1] + 1800\n", + "t_1 = times_observed_astrometry_gaia[-1]+1800\n", "\n", "# add roman observing times\n", "times_observed_astrometry_roman = []\n", - "for ii in range(t_1, t_1 + (10 * 365)):\n", + "for ii in range(t_1, t_1+(10*365)):\n", " if ii % 90 == 0:\n", - " times_observed_astrometry_roman.append(ii)\n", + " times_observed_astrometry_roman.append(ii) \n", + "\n", "\n", + "sigma_ra_roman = 5e-6 #5 micro-as\n", + "sigma_dec_roman = 5e-6 #5 micro-as\n", "\n", - "sigma_ra_roman = 5e-6 # 5 micro-as\n", - "sigma_dec_roman = 5e-6 # 5 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets,\n", - " sigma_rv,\n", + " n_planets, \n", + " sigma_rv, \n", " sigma_ra_roman,\n", " sigma_dec_roman,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv=times_observed_rv,\n", - " times_observed_astrometry=times_observed_astrometry_roman,\n", - ")\n", + " times_observed_rv = times_observed_rv,\n", + " times_observed_astrometry = times_observed_astrometry_roman\n", + " )\n", "\n", "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", "\n", - "times_observed_astrometry = np.append(\n", - " times_observed_astrometry, times[3], axis=0\n", - ")\n", + "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", "\n", "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", @@ -537,7 +532,7 @@ "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", "\n", "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", - "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)" + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" ] }, { @@ -557,40 +552,36 @@ "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib\n", - "\n", - "matplotlib.rc(\"font\", **{\"family\": \"serif\", \"serif\": [\"Computer Modern\"]})\n", - "matplotlib.rc(\"text\", usetex=True)\n", - "\n", + "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", + "matplotlib.rc('text', usetex=True)\n", "\n", "def plot_rv_signal(\n", - " n_planets,\n", - " rv_orbit,\n", - " rv_orbit_sum,\n", - " rv_sim,\n", - " rv_sim_sum,\n", - " times_rv,\n", - " times_observed_rv,\n", - "):\n", - " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", - " fig.suptitle(\"RV Signal\", fontsize=45)\n", + " n_planets, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv):\n", + " \n", + " \n", + " \n", + " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", + " fig.suptitle(\"RV Signal\", fontsize = 45)\n", " ax1 = ax\n", + " \n", + " \n", + " \n", "\n", - " ax1.plot(\n", - " times_observed_rv,\n", - " rv_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " alpha=0.3,\n", - " )\n", + " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", "\n", - " ax1.set_xlabel(\"time [BJD]\", fontsize=27)\n", - " ax1.set_ylabel(\"RV [m/s]\", fontsize=27)\n", - " ax1.legend(fontsize=27, loc=2)\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", + " ax1.legend(fontsize = 27, loc = 2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", - "\n", + " \n", " return None" ] }, @@ -602,55 +593,53 @@ "outputs": [], "source": [ "def plot_astrometry_signal(\n", - " n_planets,\n", - " ra_orbit,\n", - " ra_orbit_sum,\n", - " ra_sim,\n", - " ra_sim_sum,\n", - " dec_orbit,\n", - " dec_orbit_sum,\n", - " dec_sim,\n", + " n_planets, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry,\n", - "):\n", - " fig, ax = plt.subplots(2, 1, figsize=[13, 10], sharex=\"col\", sharey=\"row\")\n", - " fig.suptitle(\"Astrometric Signal\", fontsize=45)\n", + " times_observed_astrometry):\n", + " \n", + " \n", + "\n", + " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", + " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", " ax2 = ax[0]\n", " ax3 = ax[1]\n", + " \n", "\n", - " ax2.plot(\n", - " times_observed_astrometry,\n", - " ra_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " )\n", + "\n", + " \n", + " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", " for tick in ax3.get_xticklabels():\n", " tick.set_rotation(30)\n", - " ax2.legend(fontsize=27, loc=2)\n", - "\n", - " ax3.plot(\n", - " times_observed_astrometry,\n", - " dec_sim_sum,\n", - " \"o\",\n", - " color=\"k\",\n", - " label=\"combined signal\",\n", - " )\n", - " ax2.set_xlabel(\"time [BJD]\", fontsize=27)\n", - " ax3.legend(fontsize=27, loc=2)\n", + " ax2.legend(fontsize = 27, loc = 2)\n", + " \n", + " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", + " ax3.legend(fontsize = 27, loc = 2)\n", + "\n", + "\n", "\n", " fig.tight_layout()\n", " fig.show()\n", "\n", + "\n", + "\n", " #######\n", " #######\n", " #######\n", "\n", - " fig, ax = plt.subplots(1, 1, figsize=[10, 10])\n", + "\n", + " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", "\n", " # plot simulated data\n", - " ax.plot(ra_sim_sum, dec_sim_sum, \"o\", color=\"k\", label=\"combined signal\")\n", + " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", "\n", " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", @@ -660,7 +649,7 @@ " ax.set_title(\"initial orbit\")\n", " ax.legend()\n", " fig.show()\n", - "\n", + " \n", " return None" ] }, @@ -676,46 +665,46 @@ "cell_type": "code", "execution_count": null, "id": "82e18e64", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "plot_rv_signal(\n", - " 2, rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum, times_rv, times_observed_rv\n", - ")" + " 2, \n", + " rv_orbit, \n", + " rv_orbit_sum, \n", + " rv_sim, \n", + " rv_sim_sum, \n", + " times_rv, \n", + " times_observed_rv)\n", + " \n", + " \n" ] }, { "cell_type": "code", "execution_count": null, "id": "623b9c2d", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "plot_astrometry_signal(\n", - " 2,\n", - " ra_orbit,\n", - " ra_orbit_sum,\n", - " ra_sim,\n", - " ra_sim_sum,\n", - " dec_orbit,\n", - " dec_orbit_sum,\n", - " dec_sim,\n", + " 2, \n", + " ra_orbit, \n", + " ra_orbit_sum, \n", + " ra_sim, \n", + " ra_sim_sum, \n", + " dec_orbit, \n", + " dec_orbit_sum, \n", + " dec_sim, \n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry,\n", - ")" + " times_observed_astrometry)" ] }, { "cell_type": "markdown", "id": "99110cc0", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## Ok, now we have a simulated dataset -- let's model it!\n", "\n", @@ -736,7 +725,9 @@ " print(\"------------\")\n", "\n", " with pm.Model() as model:\n", - " # log normal prior on period around estimates\n", + " \n", + " \n", + " #log normal prior on period around estimates\n", " logP = pm.Uniform(\n", " \"logP\",\n", " lower=0,\n", @@ -744,33 +735,37 @@ " shape=2,\n", " testval=np.log(periods),\n", " )\n", + " \n", + "\n", "\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", "\n", + "\n", " ## wide uniform prior on t_periastron\n", " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", - "\n", + " \n", + " \n", " # Wide normal prior for semi-amplitude\n", - " logK = pm.Uniform(\n", - " \"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks)\n", - " )\n", - "\n", + " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", + " \n", " K = pm.Deterministic(\"K\", tt.exp(logK))\n", - "\n", + " \n", + " \n", " # Eccentricity & argument of periasteron\n", " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", "\n", + "\n", + "\n", + " \n", " # Jitter & a quadratic RV trend\n", - " # logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", "\n", " # Then we define the orbit\n", - " orbit = xo.orbits.KeplerianOrbit(\n", - " period=P, t_periastron=tperi, ecc=ecc, omega=omega\n", - " )\n", + " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", @@ -781,6 +776,8 @@ " # Sum over planets and add the background to get the full model\n", " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", + " \n", + "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", "\n", @@ -789,17 +786,19 @@ "\n", " # Finally add in the observation model. This next line adds a new contribution\n", " # to the log probability of the PyMC3 model\n", - " # err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", " err = y_rv_err\n", " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", "\n", + "\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", " map_soln = pmx.optimize(start=map_soln)\n", "\n", - " # return the max a-posteriori solution\n", + "\n", + " #return the max a-posteriori solution\n", " return map_soln" ] }, @@ -807,18 +806,16 @@ "cell_type": "code", "execution_count": null, "id": "4d6faa40", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import pymc3 as pm\n", "import pymc3_ext as pmx\n", "\n", - "np.random.seed(979)\n", + "np.random.seed(979) \n", "\n", "\n", - "# rename variables in more consistent way for modeling\n", + "#rename variables in more consistent way for modeling\n", "x_rv = np.array(times_observed_rv)\n", "y_rv = rv_sim_sum\n", "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", @@ -832,29 +829,25 @@ "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", "\n", "\n", + " \n", + " \n", "# make a fine grid that spans the observation window for plotting purposes\n", - "t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - ")\n", + "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", "# for predicted orbits\n", - "t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - ")\n", + "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", "\n", "\n", - "# minimize on RV data\n", + "#minimize on RV data\n", "periods_real = [4327.0, 365.0]\n", "periods_guess = []\n", "\n", - "# let's assume that THE can recover the period within 10% of the true value\n", + "#let's assume that THE can recover the period within 10% of the true value\n", "for per in periods_real:\n", - " periods_guess.append(\n", - " per + np.random.uniform(low=-0.1 * per, high=0.1 * per)\n", - " )\n", + " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", "\n", - "print(\"periods used as starting point = \", periods_guess)\n", + "print('periods used as starting point = ', periods_guess)\n", "periods_guess = np.array(periods_guess)\n", "\n", "\n", @@ -866,9 +859,7 @@ { "cell_type": "markdown", "id": "efc4dcc2", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## Now let's define out joint astrometry and radial velocity model:" ] @@ -880,40 +871,25 @@ "metadata": {}, "outputs": [], "source": [ - "def minimize_both(\n", - " rv_map_soln,\n", - " x_rv,\n", - " y_rv,\n", - " y_rv_err,\n", - " x_astrometry,\n", - " ra_data,\n", - " ra_err,\n", - " dec_data,\n", - " dec_err,\n", - " parallax,\n", - "):\n", - " m_sun = 333030 # earth masses\n", - "\n", - " P_RV = np.array(rv_map_soln[\"P\"])\n", - " K_RV = np.array(rv_map_soln[\"K\"])\n", - " tperi_RV = np.array(rv_map_soln[\"tperi\"])\n", - " ecc_RV = np.array(rv_map_soln[\"ecc\"])\n", - " omega_RV = np.array(rv_map_soln[\"omega\"])\n", - " min_masses_RV = (\n", - " xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value * 317.83\n", - " ) # in m_earth\n", + "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", + " m_sun = 333030 #earth masses\n", + " \n", + " P_RV = np.array(rv_map_soln['P'])\n", + " K_RV = np.array(rv_map_soln['K'])\n", + " tperi_RV = np.array(rv_map_soln['tperi'])\n", + " ecc_RV = np.array(rv_map_soln['ecc'])\n", + " omega_RV = np.array(rv_map_soln['omega'])\n", + " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", " phase_RV = determine_phase(P_RV, tperi_RV)\n", - "\n", + " \n", + " \n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - " )\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "\n", "\n", " print(\"RV Solutions\")\n", " print(\"------------\")\n", @@ -923,106 +899,98 @@ " print(\"eccentricity: \", ecc_RV)\n", " print(\"omega: \", omega_RV)\n", "\n", - " print(\"\")\n", + " print('')\n", " print(\"minimizing joint model solutions pre-MCMC\")\n", " print(\"------------\")\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", "\n", - " inc_test_vals = np.array(np.radians([5.0, 25.0, 45.0, 65.0, 85.0]))\n", + " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", " model, map_soln, logp = [], [], []\n", " for inc in inc_test_vals:\n", - " mass_test_vals = min_masses_RV / np.sin(inc)\n", - " print(\"\")\n", + " mass_test_vals = min_masses_RV/np.sin(inc)\n", + " print('')\n", " print(\"trying inclination = \" + str(np.degrees(inc)))\n", " print(\"mass test val = \" + str(mass_test_vals))\n", " print(\"------------\")\n", + " \n", + " \n", + "\n", + "\n", "\n", " def get_model():\n", " with pm.Model() as model:\n", + "\n", + "\n", " # Below we will run a version of this model where a measurement of parallax is provided\n", " # The measurement is in milliarcsec\n", " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", - " \"m_plx\", mu=parallax * 1000, sd=10, testval=parallax * 1000\n", + " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", " )\n", " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", "\n", + "\n", " # We expect the period to be around that found from just the RVs,\n", " # so we'll set a broad prior on logP\n", - "\n", + " \n", " logP = pm.Uniform(\n", - " \"logP\",\n", - " lower=0,\n", - " upper=np.log(2 * P_RV),\n", - " testval=np.log(P_RV),\n", - " shape=2,\n", + " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", " )\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", - "\n", + " \n", " # Eccentricity & argument of periasteron\n", - " ecs = pmx.UnitDisk(\n", - " \"ecs\",\n", - " shape=(2, 2),\n", - " testval=np.array(\n", - " [\n", - " np.sqrt(ecc_RV) * np.cos(omega_RV),\n", - " np.sqrt(ecc_RV) * np.sin(omega_RV),\n", - " ]\n", - " ),\n", - " )\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", + " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", + " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", + " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", - " xo.eccentricity.kipping13(\n", - " \"ecc_prior\", fixed=True, observed=ecc\n", - " )\n", + " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", + "\n", + " \n", + " \n", "\n", " # Omegas are co-dependent, so sample them with variables Omega_plus\n", - " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and\n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", " # Omega_minus is (Omega_0 - Omega_1)/2\n", - "\n", + " \n", " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", + " \n", + " \n", + " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", + " (Omega_plus - Omega_minus)] )\n", + " \n", "\n", - " Omega = tt.concatenate(\n", - " [(Omega_plus + Omega_minus), (Omega_plus - Omega_minus)]\n", - " )\n", + " Omega = pm.Deterministic(\"Omega\", Omega) \n", + " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", + " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", + " \n", "\n", - " Omega = pm.Deterministic(\"Omega\", Omega)\n", - " Omega_sum = pm.Deterministic(\n", - " \"Omega_sum\", ((Omega_plus) * 2) % np.pi\n", - " )\n", - " Omega_diff = pm.Deterministic(\n", - " \"Omega_diff\", ((Omega_minus) * 2) % np.pi\n", - " )\n", "\n", + " \n", " # For these orbits, it can also be better to fit for a phase angle\n", " # (relative to a reference time) instead of the time of periasteron\n", " # passage directly\n", " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", + " \n", "\n", + " \n", " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", - " log_m = pm.Uniform(\n", - " \"log_m\",\n", - " lower=-1,\n", - " upper=np.log(100 * mass_test_vals),\n", - " testval=np.log(mass_test_vals),\n", - " shape=2,\n", - " )\n", + " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", - " m_planet_fit = pm.Deterministic(\n", - " \"m_planet_fit\", m_planet / m_sun\n", - " )\n", + " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", "\n", - " cos_incl = pm.Uniform(\n", - " \"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2\n", - " )\n", + "\n", + " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", + " \n", + "\n", "\n", + " \n", + " \n", " # Set up the orbit\n", " orbit = xo.orbits.KeplerianOrbit(\n", " t_periastron=tperi,\n", @@ -1031,10 +999,13 @@ " ecc=ecc,\n", " omega=omega,\n", " Omega=Omega,\n", - " m_planet=m_planet_fit,\n", - " plx=plx,\n", + " m_planet = m_planet_fit,\n", + " plx=plx\n", " )\n", "\n", + "\n", + " \n", + " \n", " # Add a function for computing the full astrometry model\n", " def get_astrometry_model(t, name=\"\"):\n", " # First the astrometry induced by the planets\n", @@ -1042,56 +1013,54 @@ " # determine and print the star position at desired times\n", " pos = orbit.get_star_position(t, plx)\n", "\n", - " x, y, z = pos\n", - "\n", - " # calculate rho and theta\n", - " rhos = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", - " thetas = tt.squeeze(\n", - " tt.arctan2(y, x)\n", - " ) # radians between [-pi, pi]\n", + " x,y,z = pos\n", "\n", - " # rhos, thetas = get_star_relative_angles(t, plx)\n", - "\n", - " dec = pm.Deterministic(\n", - " \"dec\" + name, rhos * np.cos(thetas)\n", - " ) # X is north\n", - " ra = pm.Deterministic(\n", - " \"ra\" + name, rhos * np.sin(thetas)\n", - " ) # Y is east\n", "\n", + " # calculate rho and theta\n", + " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", + " \n", + " \n", + " #rhos, thetas = get_star_relative_angles(t, plx)\n", + " \n", + " \n", + " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", + " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", + " \n", " # Sum over planets to get the full model\n", - " dec_model = pm.Deterministic(\n", - " \"dec_model\" + name, tt.sum(dec, axis=-1)\n", - " )\n", - " ra_model = pm.Deterministic(\n", - " \"ra_model\" + name, tt.sum(ra, axis=-1)\n", - " )\n", + " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", + " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", + " \n", "\n", + " \n", " return dec_model, ra_model\n", "\n", + " \n", " # Define the astrometry model at the observed times\n", " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", "\n", " # Also define the model on a fine grid as computed above (for plotting)\n", - " dec_model_fine, ra_model_fine = get_astrometry_model(\n", - " t_fine, name=\"_fine\"\n", - " )\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", + "\n", + " \n", + "\n", "\n", - " # dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", - " # ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", " dec_tot_err = dec_err\n", " ra_tot_err = ra_err\n", "\n", - " # define the likelihood function, e.g., a Gaussian on both ra and dec\n", - " pm.Normal(\n", - " \"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data\n", - " )\n", - " pm.Normal(\n", - " \"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data\n", - " )\n", "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", + " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", + " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", + "\n", + "\n", + " \n", + " \n", " # ADD RV MODEL\n", "\n", + "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", " # First the RVs induced by the planets\n", @@ -1099,9 +1068,7 @@ " pm.Deterministic(\"vrad\" + name, vrad)\n", "\n", " # Sum over planets to get the full model\n", - " return pm.Deterministic(\n", - " \"rv_model\" + name, tt.sum(vrad, axis=-1)\n", - " )\n", + " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", @@ -1117,20 +1084,24 @@ " # Optimize to find the initial parameters\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", - " # map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", "\n", " map_soln = pmx.optimize(map_soln)\n", "\n", + "\n", " return model, map_soln\n", "\n", + "\n", + "\n", " a_model, a_map_soln = get_model()\n", - " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis=0)\n", - " print(\"log likelihood = \" + str(a_logp))\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", + " print('log likelihood = ' + str(a_logp))\n", "\n", " model.append(a_model)\n", " map_soln.append(a_map_soln)\n", " logp.append(a_logp)\n", "\n", + " \n", " best_index = 0\n", " for index in range(0, len(model)):\n", " if logp[index] >= logp[best_index]:\n", @@ -1140,26 +1111,31 @@ " the_map_soln = map_soln[best_index]\n", " the_logp = logp[best_index]\n", "\n", + "\n", + "\n", " return the_model, the_map_soln, the_logp\n", "\n", "\n", + "\n", "def determine_phase(P, t_periastron):\n", " phase = (2 * np.pi * t_periastron) / P\n", " return phase\n", "\n", "\n", + "\n", "def model_both(model, map_soln, tune_steps, draw_steps):\n", - " print(\"Joint RV + Astometry Minimization Solutions:\")\n", + " print('Joint RV + Astometry Minimization Solutions:')\n", " print(\"------------\")\n", "\n", - " print(\"m_planet: \" + str(map_soln[\"m_planet\"]))\n", - " print(\"P: \" + str(map_soln[\"P\"]))\n", - " print(\"incl: \" + str(map_soln[\"incl\"]))\n", - " print(\"Omega: \" + str(map_soln[\"Omega\"]))\n", - " print(\"tperi: \" + str(map_soln[\"tperi\"]))\n", - " print(\"ecc: \" + str(map_soln[\"ecc\"]))\n", - " print(\"omega: \" + str(map_soln[\"omega\"]))\n", - " print(\"plx: \" + str(map_soln[\"plx\"]))\n", + " print('m_planet: ' + str(map_soln['m_planet']))\n", + " print('P: ' + str(map_soln['P']))\n", + " print('incl: ' + str(map_soln['incl']))\n", + " print('Omega: ' + str(map_soln['Omega']))\n", + " print('tperi: ' + str(map_soln['tperi']))\n", + " print('ecc: ' + str(map_soln['ecc']))\n", + " print('omega: ' + str(map_soln['omega']))\n", + " print('plx: ' + str(map_soln['plx']))\n", + "\n", "\n", " with model:\n", " trace = pmx.sample(\n", @@ -1187,24 +1163,18 @@ "\n", "################\n", "################\n", - "# minimize on joint model\n", + "#minimize on joint model\n", "joint_model, joint_map_soln, joint_logp = minimize_both(\n", - " rv_map_soln,\n", - " x_rv,\n", - " y_rv,\n", - " y_rv_err,\n", - " x_astrometry,\n", - " ra_data,\n", - " ra_err,\n", - " dec_data,\n", - " dec_err,\n", - " parallax,\n", + " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", + " ra_data, ra_err, dec_data, dec_err, parallax\n", ")\n", "\n", "\n", + "\n", + "\n", "################\n", "################\n", - "# run full MCMC\n", + "#run full MCMC\n", "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" ] }, @@ -1223,168 +1193,143 @@ "metadata": {}, "outputs": [], "source": [ - "matplotlib.rc(\"text\", usetex=False)\n", - "\n", + "matplotlib.rc('text', usetex=False)\n", "\n", "def make_plots(trace, orbit_params, n_planets):\n", - " m_sun = 333030 # earth masses\n", + " m_sun = 333030 #earth masses\n", "\n", " [orbit_params_earth, orbit_params_jup] = orbit_params\n", "\n", - " [\n", - " P_earth,\n", - " e_earth,\n", - " Tper_earth,\n", - " omega_earth,\n", - " Omega_earth,\n", - " inclination_earth,\n", - " m_earth,\n", - " ] = orbit_params_earth\n", - "\n", - " [\n", - " P_jup,\n", - " e_jup,\n", - " Tper_jup,\n", - " omega_jup,\n", - " Omega_jup,\n", - " inclination_jup,\n", - " m_jup,\n", - " ] = orbit_params_jup\n", - "\n", - " a_true_earth = a_from_Kepler3(P_earth, 1.0 + m_earth)\n", - " a_true_jup = a_from_Kepler3(P_jup, 1.0 + m_jup)\n", - "\n", - " # [P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", - " truths = [\n", - " P_jup,\n", - " P_earth,\n", - " e_jup,\n", - " e_earth,\n", - " omega_jup,\n", - " omega_earth,\n", - " Omega_earth - Omega_jup,\n", - " Omega_earth + Omega_jup,\n", - " inclination_jup,\n", - " inclination_earth,\n", - " 0.1,\n", - " m_jup * m_sun,\n", - " m_earth * m_sun,\n", - " Tper_jup,\n", - " Tper_earth,\n", - " ]\n", - "\n", - " # [P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", - " truth_chain_plot = [\n", - " P_jup,\n", - " P_earth,\n", - " e_jup,\n", - " e_earth,\n", - " omega_jup,\n", - " omega_earth,\n", - " inclination_jup,\n", - " inclination_earth,\n", - " Tper_jup,\n", - " Tper_earth,\n", - " m_jup * m_sun,\n", - " m_earth * m_sun,\n", - " ]\n", + "\n", + " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", + " inclination_earth, m_earth] = orbit_params_earth\n", + "\n", + "\n", + " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", + " inclination_jup, m_jup] = orbit_params_jup\n", + "\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", + " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", + " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", + "\n", + "\n", + " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", + " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " # plot the table summarizing MCMC results\n", " az.summary(\n", " trace,\n", - " var_names=[\n", - " \"P\",\n", - " \"tperi\",\n", - " \"omega\",\n", - " \"Omega_sum\",\n", - " \"Omega_plus\",\n", - " \"Omega_minus\",\n", - " \"incl\",\n", - " \"ecc\",\n", - " \"plx\",\n", - " \"m_planet\",\n", - " ],\n", + " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", + " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", " )\n", "\n", + "\n", " plt.show()\n", "\n", + "\n", + "\n", + "\n", + "\n", " # plot the corner plots\n", " _ = corner.corner(\n", - " trace,\n", - " var_names=[\n", - " \"P\",\n", - " \"ecc\",\n", - " \"omega\",\n", - " \"Omega_sum\",\n", - " \"Omega_minus\",\n", - " \"incl\",\n", - " \"plx\",\n", - " \"m_planet\",\n", - " \"tperi\",\n", - " ],\n", - " quantiles=[0.16, 0.5, 0.84],\n", - " show_titles=True,\n", - " title_kwargs={\"fontsize\": 13},\n", - " truths=truths,\n", - " truth_color=\"#03003a\",\n", + " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", + " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", + " truths = truths, truth_color = \"#03003a\"\n", " )\n", "\n", " plt.show()\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # plot the chains\n", " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", " for ii in range(0, len(parameters)):\n", " plot_truth = False\n", " param = parameters[ii]\n", - "\n", - " true_vals_earth = truth_chain_plot[2 * ii]\n", - " true_vals_jup = truth_chain_plot[2 * ii + 1]\n", + " \n", + " true_vals_earth = truth_chain_plot[2*ii]\n", + " true_vals_jup = truth_chain_plot[2*ii+1]\n", " plot_truth = True\n", - "\n", - " fig, ax = plt.subplots(1, 2, figsize=(15, 3))\n", + " \n", + " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", - "\n", + " \n", " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", - "\n", - " nstep = np.arange(1, len(planet1_chain1) + 1, 1)\n", - "\n", + " \n", + " \n", + " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", + " \n", + " \n", " ax[0].plot(nstep, planet1_chain1)\n", " ax[0].plot(nstep, planet1_chain2)\n", - "\n", + " \n", " if plot_truth:\n", - " ax[0].axhline(y=true_vals_earth, color=\"r\", label=\"truth\")\n", - " ax[0].set_title(\"Jupiter\", fontsize=18)\n", - " ax[0].legend(fontsize=18)\n", - "\n", + " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", + " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", + " ax[0].legend(fontsize = 18)\n", + " \n", " ax[1].plot(nstep, planet2_chain1)\n", " ax[1].plot(nstep, planet2_chain2)\n", - "\n", + " \n", " if plot_truth:\n", - " ax[1].axhline(y=true_vals_jup, color=\"r\", label=\"truth\")\n", - " ax[1].set_title(\"Earth\", fontsize=18)\n", - " ax[1].legend(fontsize=18)\n", + " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", + " ax[1].set_title(\"Earth\", fontsize = 18)\n", + " ax[1].legend(fontsize = 18)\n", "\n", - " fig.suptitle(param, fontsize=18)\n", + " fig.suptitle(param, fontsize = 18)\n", " fig.tight_layout()\n", " plt.show()\n", "\n", - " return \"plotting complete\"\n", "\n", "\n", - "def a_from_Kepler3(period, M_tot):\n", - " period = period * 86400 # days to seconds\n", "\n", - " M_tot = M_tot * M_sun.value # solar masses to kg\n", + " return \"plotting complete\"\n", + "\n", "\n", - " a3 = (((constants.G.value) * M_tot) / (4 * np.pi**2)) * (period) ** 2.0\n", "\n", - " a = a3 ** (1 / 3)\n", "\n", - " a = a * 6.68459 * 10 ** (-12.0) # meters to AU\n", "\n", - " return a # in AUs" + "def a_from_Kepler3(period, M_tot):\n", + " period = period*86400 #days to seconds\n", + " \n", + "\n", + " M_tot = M_tot*M_sun.value #solar masses to kg\n", + " \n", + " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", + " \n", + " a = a3**(1/3)\n", + " \n", + " \n", + " \n", + " a = a * 6.68459*10**(-12.) # meters to AU\n", + " \n", + " return(a) #in AUs\n", + "\n", + "\n", + "\n" ] }, { @@ -1399,14 +1344,18 @@ "cell_type": "code", "execution_count": null, "id": "cfd9e0cd", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import arviz as az\n", "\n", - "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" + "az.summary(trace, var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"m_planet\",\n", + " \"Omega\",\n", + " \"omega\",\n", + " \"incl\"])\n" ] }, { @@ -1423,14 +1372,12 @@ "cell_type": "code", "execution_count": null, "id": "6775ae19", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ "import corner\n", "\n", - "make_plots(trace, planet_params, 2)" + "make_plots(trace, planet_params, 2)\n" ] }, { @@ -1449,31 +1396,34 @@ "outputs": [], "source": [ "import pandas as pd\n", - "\n", - "# save simulated data as dataframe\n", + "#save simulated data as dataframe \n", "simulated_data_dic = {\n", " \"times_rv_observed\": x_rv,\n", " \"rv_observed\": y_rv,\n", " \"rv_err_observed\": y_rv_err,\n", + "\n", " \"times_rv_orbit\": times_rv,\n", " \"rv_orbit\": rv_orbit,\n", " \"rv_orbit_sum\": rv_orbit_sum,\n", + "\n", " \"times_astrometry_observed\": x_astrometry,\n", " \"ra_observed\": ra_data,\n", " \"dec_observed\": dec_data,\n", " \"ra_err_observed\": ra_err,\n", " \"dec_err_observed\": dec_err,\n", + "\n", " \"times_astrometry_orbit\": times_astrometry,\n", " \"ra_orbit\": ra_orbit,\n", " \"ra_orbit_sum\": ra_orbit_sum,\n", " \"dec_orbit\": dec_orbit,\n", - " \"dec_orbit_sum\": dec_orbit_sum,\n", - "}\n", + " \"dec_orbit_sum\": dec_orbit_sum\n", + " }\n", "\n", "\n", - "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient=\"index\")\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", "simulated_data = simulated_data.transpose()\n", - "simulated_data = simulated_data.fillna(value=np.nan)" + "simulated_data = simulated_data.fillna(value=np.nan)\n", + "\n" ] }, { @@ -1485,48 +1435,55 @@ "source": [ "def remove_nans(array):\n", " array = array[~pd.isnull(array)]\n", - "\n", + " \n", " return array\n", "\n", "\n", + "\n", + "\n", "def import_orbit_data(orbit_row):\n", " import math\n", - "\n", + " \n", " orbit_row = remove_nans(orbit_row)\n", " orbit_array = []\n", " for row in orbit_row:\n", " orbit_array.append([np.array(row[1]), row[0]])\n", "\n", " orbit_array = np.array(orbit_array)\n", - "\n", + " \n", " return orbit_array\n", "\n", "\n", - "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", - " import matplotlib.pyplot as plt\n", - "\n", - " # simulated RV observations\n", - " x_rv = simulated_data[\"times_rv_observed\"].values\n", - " y_rv = simulated_data[\"rv_observed\"].values\n", - " y_rv_err = simulated_data[\"rv_err_observed\"].values\n", - "\n", - " # simulated RV full orbits\n", - " x_rv_orbit = simulated_data[\"times_rv_orbit\"].values\n", - " y_rv_orbit = simulated_data[\"rv_orbit\"].values\n", "\n", - " # simulated astrometric observations\n", - " x_astrometry = simulated_data[\"times_astrometry_observed\"].values\n", - " y_ra = simulated_data[\"ra_observed\"].values\n", - " y_dec = simulated_data[\"dec_observed\"].values\n", - " y_ra_err = simulated_data[\"ra_err_observed\"].values\n", - " y_dec_err = simulated_data[\"dec_err_observed\"].values\n", "\n", - " # simulated astrometric full orbits\n", - " x_astrometry_orbit = simulated_data[\"times_astrometry_orbit\"].values\n", - " y_ra_orbit = simulated_data[\"ra_orbit\"].values\n", - " y_dec_orbit = simulated_data[\"dec_orbit\"].values\n", + "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", + " import matplotlib.pyplot as plt\n", "\n", - " # remove any nans from simulated data (created because of mismatching lengths)\n", + " \n", + " \n", + " #simulated RV observations\n", + " x_rv = simulated_data['times_rv_observed'].values\n", + " y_rv = simulated_data['rv_observed'].values\n", + " y_rv_err = simulated_data['rv_err_observed'].values\n", + " \n", + " #simulated RV full orbits\n", + " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", + " y_rv_orbit = simulated_data['rv_orbit'].values\n", + " \n", + " #simulated astrometric observations\n", + " x_astrometry = simulated_data['times_astrometry_observed'].values\n", + " y_ra = simulated_data['ra_observed'].values\n", + " y_dec = simulated_data['dec_observed'].values\n", + " y_ra_err = simulated_data['ra_err_observed'].values\n", + " y_dec_err = simulated_data['dec_err_observed'].values\n", + " \n", + " #simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", + " y_ra_orbit = simulated_data['ra_orbit'].values\n", + " y_dec_orbit = simulated_data['dec_orbit'].values\n", + " \n", + " \n", + " #remove any nans from simulated data (created because of mismatching lengths)\n", " x_rv = remove_nans(x_rv)\n", " y_rv = remove_nans(y_rv)\n", " y_rv_err = remove_nans(y_rv_err)\n", @@ -1540,34 +1497,43 @@ " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", - "\n", + " \n", + " \n", + " \n", + " \n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(\n", - " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", - " )\n", + " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(\n", - " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", - " )\n", - "\n", + " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " \n", + " \n", + " \n", + " \n", + " \n", " # add NaN for plotting purpose on true model\n", - " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float(\"NaN\"))\n", - " y_ra_orbit = np.insert(y_ra_orbit, 10000, float(\"NaN\"), axis=0)\n", - " y_dec_orbit = np.insert(y_dec_orbit, 10000, float(\"NaN\"), axis=0)\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", + " \n", + " \n", + " \n", + "\n", "\n", + " \n", + " \n", " figsAs = []\n", " figsRV = []\n", " for n in range(0, nplanets):\n", + " \n", " ######\n", - " # make rv plots\n", - "\n", - " figRV, ax = plt.subplots(\n", - " 2, 1, figsize=[13, 6], gridspec_kw={\"height_ratios\": [3, 1]}\n", - " )\n", + " #make rv plots\n", + " \n", + " \n", + " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", " ax0, ax1 = ax[0], ax[1]\n", - "\n", + " \n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1578,68 +1544,86 @@ " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # Plot the data\n", - " ax0.errorbar(\n", - " x_rv,\n", - " y_rv - other,\n", - " yerr=y_rv_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - "\n", + " \n", " med_rv = []\n", " for an_x in x_rv:\n", " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", "\n", + "\n", + "\n", + " \n", " # plot the MCMC model\n", " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(t_rv, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax0.fill_between(\n", + " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", - "\n", + " \n", " # plot the true RA model\n", - " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color=\"k\", label=\"truth\")\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", - "\n", - " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " # plot residuals\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " #plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_rv,\n", - " y_rv - other - med_rv,\n", - " yerr=y_rv_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_rv, y_rv - other - med_rv,\n", + " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", " )\n", - "\n", + " \n", + " \n", " figRV.tight_layout()\n", " figRV.subplots_adjust(hspace=0.5)\n", - " figRV.suptitle(\"planet {0}\".format(n + 1), fontsize=33)\n", + " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", " figRV.tight_layout()\n", - "\n", + " \n", + " \n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", " figRV.show()\n", " ################\n", - " # plot astrometry\n", - "\n", - " figAs, ax = plt.subplots(\n", - " 2,\n", - " 2,\n", - " figsize=[18, 9],\n", - " sharey=\"row\",\n", - " gridspec_kw={\"height_ratios\": [3, 1]},\n", - " )\n", + " #plot astrometry\n", + " \n", + " \n", + "\n", + " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", "\n", " # Get the posterior median orbital parameters\n", @@ -1652,16 +1636,14 @@ " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", + "\n", + "\n", + "\n", + "\n", + "\n", " # plot the DEC data\n", - " ax0.errorbar(\n", - " x_astrometry,\n", - " y_dec - other,\n", - " yerr=y_dec_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", "\n", " # plot the MCMC DEC model\n", " pred = np.percentile(\n", @@ -1669,78 +1651,55 @@ " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - "\n", + " \n", " med_dec = []\n", " for an_x in x_astrometry:\n", " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", + " \n", + "\n", "\n", " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax0.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", "\n", " # plot the true DEC model\n", - " ax0.plot(\n", - " x_astrometry_orbit, y_dec_orbit.T[n], color=\"k\", label=\"truth\"\n", - " )\n", + " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + " \n", + " \n", + " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", + " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", + " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + "\n", + "\n", "\n", - " ax0.arrow(\n", - " 0.5,\n", - " 0.0,\n", - " 0,\n", - " 1,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.0,\n", - " head_length=0.0,\n", - " color=\"k\",\n", - " ls=\"--\",\n", - " )\n", - " ax0.text(\n", - " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax0.transAxes\n", - " )\n", - " ax0.text(\n", - " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax0.transAxes\n", - " )\n", - " ax0.arrow(\n", - " 0.49,\n", - " 0.87,\n", - " -0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", - " ax0.arrow(\n", - " 0.51,\n", - " 0.87,\n", - " 0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax0.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize=18)\n", - "\n", - " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", + " \n", + " \n", + " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - "\n", - " # plot residuals\n", + " \n", + " \n", + " #plot residuals \n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_astrometry,\n", - " y_dec - other - med_dec,\n", - " yerr=y_dec_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_astrometry, y_dec - other - med_dec,\n", + " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", " )\n", "\n", + "\n", + "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1750,99 +1709,86 @@ " other = np.median(\n", " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", + " \n", + "\n", + "\n", + "\n", + "\n", "\n", " # Plot the RA data\n", - " ax2.errorbar(\n", - " x_astrometry,\n", - " y_ra - other,\n", - " yerr=y_ra_err,\n", - " fmt=\".\",\n", - " color=\"#00257c\",\n", - " label=\"data\",\n", - " alpha=0.7,\n", - " )\n", + " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", + " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + "\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", " [16, 50, 84],\n", - " axis=(0, 1),\n", + " axis=(0, 1), \n", " )\n", - "\n", + " \n", + " \n", " med_ra = []\n", " for an_x in x_astrometry:\n", " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", "\n", + "\n", " # plot the MCMC RA model\n", " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax2.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", + " art = ax2.fill_between(\n", + " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", + " )\n", " art.set_edgecolor(\"none\")\n", - "\n", + " \n", " # plot the true RA model\n", - " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color=\"k\", label=\"truth\")\n", - "\n", - " ax2.arrow(\n", - " 0.5,\n", - " 0.0,\n", - " 0,\n", - " 1,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.0,\n", - " head_length=0.0,\n", - " color=\"k\",\n", - " ls=\"--\",\n", - " )\n", - " ax2.text(\n", - " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax2.transAxes\n", - " )\n", - " ax2.text(\n", - " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax2.transAxes\n", - " )\n", - " ax2.arrow(\n", - " 0.49,\n", - " 0.87,\n", - " -0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", - " ax2.arrow(\n", - " 0.51,\n", - " 0.87,\n", - " 0.15,\n", - " 0,\n", - " width=0.001,\n", - " transform=ax2.transAxes,\n", - " head_width=0.03,\n", - " head_length=0.03,\n", - " color=\"k\",\n", - " )\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", + "\n", + "\n", + "\n", + " \n", + " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", + " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", + " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", + " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", + " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", + " head_width= 0.03, head_length = .03, color ='k')\n", "\n", + "\n", + " \n", " ax2.legend(fontsize=9, loc=1)\n", " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize=18)\n", - "\n", - " ax3.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", + " \n", + " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " # plot residuals\n", + " \n", + " #plot residuals\n", " ax3.axhline(0.0, color=\"k\")\n", " ax3.errorbar(\n", - " x_astrometry,\n", - " y_ra - other - med_ra,\n", - " yerr=y_ra_err,\n", - " color=\"#00257c\",\n", - " fmt=\".\",\n", + " x_astrometry, y_ra - other - med_ra,\n", + " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", " )\n", - "\n", + " \n", " figAs.tight_layout()\n", - " figAs.subplots_adjust(wspace=0.2)\n", - "\n", + " figAs.subplots_adjust(wspace=.2)\n", + " \n", + " \n", " figAs.show()\n", - "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "\n", + " \n", " return None" ] }, @@ -1850,20 +1796,16 @@ "cell_type": "code", "execution_count": null, "id": "b3941ebb", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "outputs": [], "source": [ - "make_rv_astrometry_plots(simulated_data, trace, 2)" + "make_rv_astrometry_plots(simulated_data, trace, 2)\n" ] }, { "cell_type": "markdown", "id": "1f6c90cd", - "metadata": { - "lines_to_next_cell": 2 - }, + "metadata": {}, "source": [ "## And lastly, check how well our MCMC posterior median +/- 1 sigma results agree with the simulated values" ] @@ -1876,7 +1818,9 @@ "outputs": [], "source": [ "def load_posterior_params(trace, nplanets):\n", - " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\n", + "\n", + " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", + "\n", "\n", " params_earth = defaultdict(list)\n", " params_jup = defaultdict(list)\n", @@ -1891,139 +1835,138 @@ " posterior_meds = defaultdict(list)\n", " posterior_errs = defaultdict(list)\n", " for param in parameters:\n", + "\n", + "\n", " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", "\n", - " planet_quantile = [\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.84),\n", - " ]\n", + " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", + "\n", + "\n", "\n", " planet_err = [\n", " planet_med - planet_quantile[0],\n", - " planet_quantile[1] - planet_med,\n", + " planet_quantile[1] - planet_med \n", " ]\n", "\n", - " # print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + "\n", + "\n", + " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", "\n", " posterior_meds[param] = np.round(planet_med, 3)\n", " posterior_errs[param] = np.round(planet_err, 3)\n", - " table_params[param].append(\n", - " str(np.round(planet_med, 3))\n", - " + \" (+\"\n", - " + str(np.round(planet_err[1], 3))\n", - " + \" -\"\n", - " + str(np.round(planet_err[0], 3))\n", - " + \")\"\n", - " )\n", + " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", + " str(np.round(planet_err[1], 3)) + \" -\" + \n", + " str(np.round(planet_err[0], 3)) + \")\")\n", + "\n", + "\n", + "\n", + "\n", + "\n", "\n", " table_params_all.append(table_params)\n", " posterior_meds_all.append(posterior_meds)\n", " posterior_errs_all.append(posterior_errs)\n", - " # print(\"\")\n", - " # print(\"\")\n", + " #print(\"\")\n", + " #print(\"\")\n", + "\n", + "\n", "\n", " return table_params_all, posterior_meds_all, posterior_errs_all\n", "\n", "\n", - "def make_comparison_figures(\n", - " input_params,\n", - " posterior_orbit_params,\n", - " keys_orbit_params,\n", - " posterior_meds,\n", - " posterior_errs,\n", - " nplanets,\n", - "):\n", + "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", + " posterior_meds, posterior_errs, nplanets):\n", + "\n", " import matplotlib\n", " import matplotlib.pyplot as plt\n", - "\n", + " \n", + " \n", " input_params = np.array(input_params)\n", - " input_params.T[-1] = input_params.T[-1] * m_sun\n", + " input_params.T[-1] = input_params.T[-1]*m_sun\n", "\n", " comparisons = []\n", " for ii in range(0, nplanets):\n", - " fig, axs = plt.subplots(\n", - " 1, 2, figsize=[13, 6], gridspec_kw={\"width_ratios\": [1, 1]}\n", - " )\n", + " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", " plt.subplots_adjust(hspace=1)\n", - "\n", + " \n", " #### add comparison plots\n", - "\n", + " \n", + " \n", " input_mass = input_params[ii][-1]\n", " input_period = input_params[ii][0]\n", + " \n", + " xs = posterior_meds[ii]['m_planet']\n", + " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", "\n", - " xs = posterior_meds[ii][\"m_planet\"]\n", - " xs_err = np.array([posterior_errs[ii][\"m_planet\"]]).T\n", - "\n", - " ys = posterior_meds[ii][\"P\"]\n", - " ys_err = np.array([posterior_errs[ii][\"P\"]]).T\n", - "\n", - " color = \"#1c245f\"\n", - "\n", - " axs[0].errorbar(\n", - " xs,\n", - " ys,\n", - " xerr=xs_err,\n", - " yerr=ys_err,\n", - " marker=\"o\",\n", - " color=color,\n", - " markersize=13,\n", - " label=r\"med and 1$\\sigma$ errs\",\n", - " linewidth=3,\n", - " alpha=0.7,\n", - " )\n", + " ys = posterior_meds[ii]['P']\n", + " ys_err = np.array([posterior_errs[ii]['P']]).T\n", + "\n", + "\n", + "\n", + " color = '#1c245f'\n", + "\n", + " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", + " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", + " linewidth = 3, alpha = 0.7)\n", + "\n", + " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", "\n", - " axs[0].axvline(\n", - " x=input_mass, ymin=0, ymax=1, color=\"k\", ls=\"--\", label=\"simulated\"\n", - " )\n", - " axs[0].axhline(y=input_period, xmin=0, xmax=1, color=\"k\", ls=\"--\")\n", "\n", - " axs[0].set_xlabel(r\"mass [M$_{earth}$]\", fontsize=27)\n", - " axs[0].set_ylabel(r\"period [days]\", fontsize=27)\n", + " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", + " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", "\n", - " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", "\n", - " axs[0].legend(fontsize=13)\n", + " fig.suptitle('input and posterior parameters', fontsize = 27)\n", "\n", + " axs[0].legend(fontsize = 13)\n", + " \n", + " \n", + " \n", " #### add comparison tables\n", "\n", + "\n", " sim_params = input_params[ii]\n", " post_params = posterior_orbit_params[ii]\n", - " data = np.column_stack(\n", - " (\n", - " np.array(keys_orbit_params),\n", - " np.round(np.array(sim_params), 3),\n", - " np.array(list(post_params.values())),\n", - " )\n", - " )\n", + " data = np.column_stack((np.array(keys_orbit_params), \n", + " np.round(np.array(sim_params),3),\n", + " np.array(list(post_params.values()))))\n", "\n", - " labels = [\"param\", \"truth\", r\"MCMC med +/- 1$\\sigma$\"]\n", + " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", "\n", " df = pd.DataFrame(data, columns=labels)\n", "\n", - " table = axs[1].table(\n", - " cellText=df.values,\n", - " colLabels=df.columns,\n", - " colColours=[\"#D7E5F0\"] * 3,\n", - " loc=\"center\",\n", - " bbox=[0, 0, 1, 1],\n", - " colWidths=[0.2, 0.2, 0.6],\n", - " )\n", - " # for (row, col), cell in table.get_celld().items():\n", + " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", + " colColours =[\"#D7E5F0\"]*3, loc='center', \n", + " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", + " #for (row, col), cell in table.get_celld().items():\n", " # if (row == 0):\n", " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", "\n", + "\n", " table.auto_set_column_width(col=labels)\n", "\n", - " # ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", - " # t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", - " # t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", " axs[1].set_axis_off()\n", "\n", + "\n", " table.auto_set_font_size(False)\n", " table.set_fontsize(13)\n", "\n", - " comparisons.append(fig)\n", + " \n", + " \n", + " \n", + "\n", + "\n", "\n", + " comparisons.append(fig)\n", + " \n", + " \n", + " \n", + " \n", " return comparisons" ] }, @@ -2037,29 +1980,12 @@ "from collections import defaultdict\n", "\n", "\n", - "(\n", - " table_params_all,\n", - " posterior_meds_all,\n", - " posterior_errs_all,\n", - ") = load_posterior_params(trace, 2)\n", - "\n", - "planet_param_keys = [\n", - " \"period\",\n", - " \"ecc\",\n", - " r\"T$_\\mathrm{per}$\",\n", - " r\"$\\omega$\",\n", - " r\"$\\Omega$\",\n", - " \"inc\",\n", - " \"mass\",\n", - "]\n", - "make_comparison_figures(\n", - " [planet_params[1], planet_params[0]],\n", - " table_params_all,\n", - " planet_param_keys,\n", - " posterior_meds_all,\n", - " posterior_errs_all,\n", - " 2,\n", - ")" + "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", + "load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", + "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", + " posterior_meds_all, posterior_errs_all, 2)" ] }, { From b6f96ec2659719df2e20901b2ade42d646fd0a95 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 15 Mar 2023 17:36:20 +0000 Subject: [PATCH 6/6] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- docs/tutorials/astrometry_rv_case_study.ipynb | 1740 +++++++++-------- 1 file changed, 907 insertions(+), 833 deletions(-) diff --git a/docs/tutorials/astrometry_rv_case_study.ipynb b/docs/tutorials/astrometry_rv_case_study.ipynb index af760fdd..6a29e62f 100644 --- a/docs/tutorials/astrometry_rv_case_study.ipynb +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -12,7 +12,9 @@ "cell_type": "code", "execution_count": null, "id": "4abd3961", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import exoplanet as xo\n", @@ -20,8 +22,7 @@ "\n", "\n", "xo.utils.docs_setup()\n", - "print(f\"exoplanet.__version__ = '{xo.__version__}'\")\n", - "\n" + "print(f\"exoplanet.__version__ = '{xo.__version__}'\")" ] }, { @@ -29,7 +30,7 @@ "id": "9a87ff47", "metadata": {}, "source": [ - "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023. \n", + "In this case study we’ll walk through a joint astrometric and radial velocity model with exoplanet. This example is done using all simulated data, based on the expected 1$\\sigma$ measurement error per observation of radial velocity observations with the Terra Hunting Experiment and astrometry with Gaia and the Nancy Grace Roman Space Telescope. This case study is one specific example of the full simulations/models run and presented in Yahalomi et al. 2023.\n", "\n", "The Terra Hunting Experiment (THE), which will be run using the HARPS3 instrument on the 2.5m Isaac Newton Telescope in La Palma in the Canary Islands, will observe at least 40 bright and nearby G and K dwarfs searching for exoplanets that are Earth-like in both mass and orbital period for 10 years. With a 1$\\sigma$ measurement error, and approximately 1,800 observations over a decade, Hall et al. 2018 showed that THE will be capabale of detecting Earth analogs. To read more about THE, check out their website: https://www.terrahunting.org/.\n", "\n", @@ -49,12 +50,14 @@ { "cell_type": "markdown", "id": "93ceb2a9", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## (Simulated) Data\n", "\n", "\n", - "First, we need to simulate our dataset. \n", + "First, we need to simulate our dataset.\n", "\n", "\n", "### Let's define 2 helpful functions in order to simulate our data...." @@ -68,27 +71,24 @@ "outputs": [], "source": [ "def create_orbit(n_planets, orbit_params):\n", - " \n", - " '''\n", + " \"\"\"\n", " inputs:\n", " - n_planets = number of planets\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists:\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", - " '''\n", - " \n", + " - m_planet = mass of planet in Earth masses\n", + " \"\"\"\n", "\n", - " \n", " period = []\n", " ecc = []\n", - " Tper = [] \n", - " omega = [] \n", - " Omega = [] \n", + " Tper = []\n", + " omega = []\n", + " Omega = []\n", " incl = []\n", " m_planet = []\n", " for ii in range(0, n_planets):\n", @@ -101,15 +101,15 @@ " m_planet.append(orbit_params[ii][6])\n", "\n", " orbit = xo.orbits.KeplerianOrbit(\n", - " period = period,\n", - " ecc = ecc, \n", - " t_periastron = Tper, \n", - " omega = omega, \n", - " Omega = Omega, \n", - " incl = incl,\n", - " m_planet = m_planet) \n", - " \n", - " \n", + " period=period,\n", + " ecc=ecc,\n", + " t_periastron=Tper,\n", + " omega=omega,\n", + " Omega=Omega,\n", + " incl=incl,\n", + " m_planet=m_planet,\n", + " )\n", + "\n", " return orbit" ] }, @@ -120,24 +120,24 @@ "metadata": {}, "outputs": [], "source": [ - "np.random.seed(979) \n", + "np.random.seed(979)\n", + "\n", "\n", "def simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra,\n", " sigma_dec,\n", " plx,\n", " orbit_params,\n", - " times_observed_rv = None,\n", - " times_observed_astrometry = None,\n", - " t_dur_rv = None, \n", - " n_obs_rv = None,\n", - " t_dur_astrometry = None, \n", - " n_obs_astrometry = None \n", - " ):\n", - "\n", - " '''\n", + " times_observed_rv=None,\n", + " times_observed_astrometry=None,\n", + " t_dur_rv=None,\n", + " n_obs_rv=None,\n", + " t_dur_astrometry=None,\n", + " n_obs_astrometry=None,\n", + "):\n", + " \"\"\"\n", " inputs:\n", " - orbit = keplerian orbit generated by exoplanet code\n", " - n_planets = number of planets\n", @@ -147,230 +147,231 @@ " - plx = parallax of system\n", " - orbit_params = orbital parameters for each planet arranged as a list of lists of shape (7 x n_planets):\n", " - period = period in days\n", - " - ecc = eccentricity \n", + " - ecc = eccentricity\n", " - t_periastron = The epoch of a reference periastron passage in days.\n", " - omega = argument of periapsis in radans\n", - " - Omega = longitude of ascending node \n", + " - Omega = longitude of ascending node\n", " - incl = inclination in radians\n", - " - m_planet = mass of planet in Earth masses \n", + " - m_planet = mass of planet in Earth masses\n", " - times_observed_rv = (default=None) the observed times for RV observations\n", " - times_observed_astrometry = (default=None) the observed times for astrometry observations\n", " - t_dur_rv = (default=None) duration of RV observations in days (only if we don't have a set times of observed observations)\n", " - n_obs_rv = (default=None) number of RV observations (only if we don't have a set times of observed observations)\n", " - t_dur_astrometry = (default=None) duration of astrometry observations (only if we don't have a set times of observed observations)\n", " - n_obs_astrometry = (default=None) number of astrometry observations (only if we don't have a set times of observed observations)\n", - " \n", - " \n", - " '''\n", - " \n", + "\n", + "\n", + " \"\"\"\n", + "\n", " if n_planets == 1:\n", " orbit_params = [orbit_params]\n", - " \n", - " \n", - " #if times_observed_rv != None then just define times_rv and linspace between min and max\n", - " if times_observed_rv != None:\n", - " times_rv = np.linspace(np.min(times_observed_rv), np.max(times_observed_rv), 10000)\n", - "\n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined\n", + "\n", + " # if times_observed_rv != None then just define times_rv and linspace between min and max\n", + " if times_observed_rv != None:\n", + " times_rv = np.linspace(\n", + " np.min(times_observed_rv), np.max(times_observed_rv), 10000\n", + " )\n", + "\n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_rv = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_rv:\n", " start_time_rv = orbit_params[ii][2]\n", "\n", + " # define RV observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_rv = np.linspace(start_time_rv, start_time_rv + t_dur_rv, 10000)\n", + " times_observed_rv = np.linspace(\n", + " start_time_rv, start_time_rv + t_dur_rv, n_obs_rv\n", + " )\n", "\n", - " \n", - " #define RV observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, 10000) \n", - " times_observed_rv = np.linspace(start_time_rv, start_time_rv+t_dur_rv, n_obs_rv)\n", - " \n", - "\n", + " # if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", + " if times_observed_astrometry != None:\n", + " times_astrometry = np.linspace(\n", + " np.min(times_observed_astrometry),\n", + " np.max(times_observed_astrometry),\n", + " 10000,\n", + " )\n", "\n", - " #if times_observed_astrometry != None then just define times_rv and linspace between min and max\n", - " if times_observed_astrometry != None:\n", - " times_astrometry = np.linspace(np.min(times_observed_astrometry), np.max(times_observed_astrometry), 10000)\n", - " \n", - " #define start_time_RV, which will be used to determine the times to simulate \n", - " #if times_observed_RV and times_observed_astrometry aren't defined \n", + " # define start_time_RV, which will be used to determine the times to simulate\n", + " # if times_observed_RV and times_observed_astrometry aren't defined\n", " else:\n", " start_time_astrometry = np.inf\n", " for ii in range(0, n_planets):\n", " if orbit_params[ii][2] < start_time_astrometry:\n", - " start_time_astrometry = orbit_params[ii][2] \n", - "\n", - "\n", - " #define astrometry observation time\n", - " #100000 is arbitrary large number finely spaced enough to show details in curve\n", - " times_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, 10000) \n", - " times_observed_astrometry = np.linspace(start_time_astrometry, start_time_astrometry+t_dur_astrometry, n_obs_astrometry)\n", - "\n", + " start_time_astrometry = orbit_params[ii][2]\n", + "\n", + " # define astrometry observation time\n", + " # 100000 is arbitrary large number finely spaced enough to show details in curve\n", + " times_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " 10000,\n", + " )\n", + " times_observed_astrometry = np.linspace(\n", + " start_time_astrometry,\n", + " start_time_astrometry + t_dur_astrometry,\n", + " n_obs_astrometry,\n", + " )\n", "\n", - " \n", " orbit = create_orbit(n_planets, orbit_params)\n", - " \n", + "\n", " rv = orbit.get_radial_velocity(times_rv)\n", " rv_observed = orbit.get_radial_velocity(times_observed_rv)\n", - " \n", + "\n", " rv_orbit = rv.eval()\n", "\n", " if n_planets > 1:\n", - " rv_orbit_sum = np.sum(rv_orbit, axis = 1)\n", + " rv_orbit_sum = np.sum(rv_orbit, axis=1)\n", " else:\n", " rv_orbit_sum = rv_orbit\n", - " \n", "\n", " rv_observed = rv_observed.eval()\n", " if n_planets > 1:\n", - " rv_observed_sum = np.sum(rv_observed, axis = 1)\n", + " rv_observed_sum = np.sum(rv_observed, axis=1)\n", " else:\n", " rv_observed_sum = rv_observed\n", "\n", - "\n", - " #-----------\n", - " #-----------\n", + " # -----------\n", + " # -----------\n", "\n", " # determine and print the star position at desired times\n", " pos = theano.function([], orbit.get_star_position(times_astrometry, plx))()\n", "\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", - " #pos = tt.sum(pos, axis=-1)\n", - "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rho = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", + " rho = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", " theta = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", + "\n", " rho, theta = rho.eval(), theta.eval()\n", - " \n", - " #rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", + " # rho, theta = theano.function([], get_star_relative_angles(times_astrometry, plx))()\n", "\n", " rho_orbit = rho\n", - " \n", + "\n", " if n_planets > 1:\n", - " rho_orbit_sum = np.sum(rho_orbit, axis = 1)\n", + " rho_orbit_sum = np.sum(rho_orbit, axis=1)\n", " else:\n", " rho_orbit_sum = rho_orbit\n", "\n", " theta_orbit = theta\n", - " \n", + "\n", " if n_planets > 1:\n", - " theta_orbit_sum = np.sum(theta_orbit, axis = 1)\n", + " theta_orbit_sum = np.sum(theta_orbit, axis=1)\n", " else:\n", " theta_orbit_sum = theta_orbit\n", - " \n", - " \n", + "\n", " # when summing over theta, position angle, we have to careful because position\n", - " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract \n", + " # angle has the range -pi to pi. So for only summing 2 thetas, we can subtract\n", " # 2pi whenever theta_sum > pi and add 2pi whenever theta_sum < -pi to get back\n", " # in the correct range. Be careful though if modeling more than 2 planets, this\n", " # doesn't completely solve the problem!\n", - " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2*np.pi\n", - " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2*np.pi\n", - "\n", + " theta_orbit_sum[theta_orbit_sum > np.pi] -= 2 * np.pi\n", + " theta_orbit_sum[theta_orbit_sum < -np.pi] += 2 * np.pi\n", "\n", " # determine and print the star position at desired times\n", - " pos_observed = theano.function([], orbit.get_star_position(times_observed_astrometry, plx))()\n", - "\n", - "\n", - " #pos = tt.sum(pos, axis=-1)\n", + " pos_observed = theano.function(\n", + " [], orbit.get_star_position(times_observed_astrometry, plx)\n", + " )()\n", "\n", - " x_obs,y_obs,z_obs = pos_observed\n", + " # pos = tt.sum(pos, axis=-1)\n", "\n", + " x_obs, y_obs, z_obs = pos_observed\n", "\n", " # calculate rho and theta\n", - " rho_observed = tt.squeeze(tt.sqrt(x_obs ** 2 + y_obs ** 2)) # arcsec\n", - " theta_observed = tt.squeeze(tt.arctan2(y_obs, x_obs)) # radians between [-pi, pi]\n", - " \n", - " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", + " rho_observed = tt.squeeze(tt.sqrt(x_obs**2 + y_obs**2)) # arcsec\n", + " theta_observed = tt.squeeze(\n", + " tt.arctan2(y_obs, x_obs)\n", + " ) # radians between [-pi, pi]\n", "\n", + " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\n", "\n", - " \n", " if n_planets > 1:\n", - " rho_observed_sum = np.sum(rho_observed, axis = 1)\n", + " rho_observed_sum = np.sum(rho_observed, axis=1)\n", " else:\n", " rho_observed_sum = rho_observed\n", "\n", - " \n", " if n_planets > 1:\n", - " theta_observed_sum = np.sum(theta_observed, axis = 1)\n", + " theta_observed_sum = np.sum(theta_observed, axis=1)\n", " else:\n", " theta_observed_sum = theta_observed\n", - " \n", - " \n", + "\n", " #######\n", " #######\n", " #######\n", "\n", " # convert rho and theta into ra and dec\n", - " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", + " ra_orbit = rho_orbit * np.sin(theta_orbit) # +ra is east\n", " dec_orbit = rho_orbit * np.cos(theta_orbit) # +dec is north\n", "\n", " if n_planets > 1:\n", - " ra_orbit_sum = np.sum(ra_orbit, axis = 1)\n", + " ra_orbit_sum = np.sum(ra_orbit, axis=1)\n", " else:\n", " ra_orbit_sum = ra_orbit\n", "\n", - " \n", " if n_planets > 1:\n", - " dec_orbit_sum = np.sum(dec_orbit, axis = 1)\n", + " dec_orbit_sum = np.sum(dec_orbit, axis=1)\n", " else:\n", " dec_orbit_sum = dec_orbit\n", "\n", - " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", + " ra_observed = rho_observed * np.sin(theta_observed) # +ra is east\n", " dec_observed = rho_observed * np.cos(theta_observed) # +dec is north\n", "\n", - "\n", - "\n", - "\n", - " \n", - " #----------- \n", - " #----------- \n", - " #----------- simulate data\n", - " #----------- \n", - " #----------- \n", - "\n", + " # -----------\n", + " # -----------\n", + " # ----------- simulate data\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", " rv_sim = rv_observed\n", - " rv_sim_sum = np.sum(rv_observed, axis = 1) + np.random.normal(0, sigma_rv, len(rv_observed))\n", + " rv_sim_sum = np.sum(rv_observed, axis=1) + np.random.normal(\n", + " 0, sigma_rv, len(rv_observed)\n", + " )\n", " else:\n", " rv_sim = rv_observed + np.random.normal(0, sigma_rv, len(rv_observed))\n", " rv_sim_sum = rv_sim\n", "\n", - " #-----------\n", - " #-----------\n", - "\n", + " # -----------\n", + " # -----------\n", "\n", " if n_planets > 1:\n", - " ra_sim = ra_observed \n", - " ra_sim_sum = np.sum(ra_sim, axis = 1) + np.random.normal(0, sigma_ra, len(ra_observed))\n", - " \n", + " ra_sim = ra_observed\n", + " ra_sim_sum = np.sum(ra_sim, axis=1) + np.random.normal(\n", + " 0, sigma_ra, len(ra_observed)\n", + " )\n", + "\n", " else:\n", " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", " ra_sim_sum = ra_sim\n", - " \n", - " \n", - " \n", + "\n", " if n_planets > 1:\n", - " dec_sim = dec_observed \n", - " dec_sim_sum = np.sum(dec_sim, axis = 1) + np.random.normal(0, sigma_dec, len(dec_observed))\n", - " \n", + " dec_sim = dec_observed\n", + " dec_sim_sum = np.sum(dec_sim, axis=1) + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", + "\n", " else:\n", - " dec_sim = dec_observed + np.random.normal(0, sigma_dec, len(dec_observed))\n", + " dec_sim = dec_observed + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\n", " dec_sim_sum = dec_sim\n", "\n", - " \n", - " times = [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry]\n", + " times = [\n", + " times_rv,\n", + " times_observed_rv,\n", + " times_astrometry,\n", + " times_observed_astrometry,\n", + " ]\n", " rv_results = [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum]\n", " ra_results = [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum]\n", " dec_results = [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]\n", - " \n", - " \n", + "\n", " return [times, rv_results, ra_results, dec_results]" ] }, @@ -398,10 +399,20 @@ "source": [ "import numpy as np\n", "\n", - "parallax = 0.1 #as \n", - "n_planets=2\n", - "planet_params = [[365.0,0.0167,100.0,1.7959438003021653,0.0,np.pi/4,3e-06],\n", - " [4327.0,0.0484,500.0,-1.4957471689591395,1.7523105690023069,np.pi/4+0.022759,0.000955]]" + "parallax = 0.1 # as\n", + "n_planets = 2\n", + "planet_params = [\n", + " [365.0, 0.0167, 100.0, 1.7959438003021653, 0.0, np.pi / 4, 3e-06],\n", + " [\n", + " 4327.0,\n", + " 0.0484,\n", + " 500.0,\n", + " -1.4957471689591395,\n", + " 1.7523105690023069,\n", + " np.pi / 4 + 0.022759,\n", + " 0.000955,\n", + " ],\n", + "]" ] }, { @@ -422,28 +433,28 @@ "cell_type": "code", "execution_count": null, "id": "55938572", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "from aesara_theano_fallback import aesara as theano\n", "import theano.tensor as tt\n", "\n", "\n", - "#add gaia observing times\n", + "# add gaia observing times\n", "times_observed_astrometry_gaia = []\n", "t_0 = 0\n", - "for ii in range(t_0, t_0+3600):\n", - " if ii % (3600/100) == 0:\n", + "for ii in range(t_0, t_0 + 3600):\n", + " if ii % (3600 / 100) == 0:\n", " times_observed_astrometry_gaia.append(ii)\n", "\n", "\n", - "\n", - "#add THE observing times\n", + "# add THE observing times\n", "times_observed_rv = []\n", - "t_0 = 3300 #Gaia first light (2014) ~9 years before THE first expected light (2023)\n", + "t_0 = 3300 # Gaia first light (2014) ~9 years before THE first expected light (2023)\n", "add_data = True\n", - "for ii in range(t_0, t_0+3600):\n", - "\n", + "for ii in range(t_0, t_0 + 3600):\n", " if ii % 180 == 0:\n", " if add_data:\n", " add_data = False\n", @@ -454,72 +465,66 @@ " times_observed_rv.append(ii)\n", "\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "sigma_rv = 0.3\n", "\n", - "sigma_ra_gaia = 6e-5 #60 micro-as\n", - "sigma_dec_gaia = 6e-5 #60 micro-as\n", - "\n", + "sigma_ra_gaia = 6e-5 # 60 micro-as\n", + "sigma_dec_gaia = 6e-5 # 60 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_gaia,\n", " sigma_dec_gaia,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_gaia\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_gaia,\n", + ")\n", "\n", "\n", - "[[times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", - "[rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", - "[ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", - "[dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum]] = times, rv_results, ra_results, dec_results\n", + "[\n", + " [times_rv, times_observed_rv, times_astrometry, times_observed_astrometry],\n", + " [rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum],\n", + " [ra_orbit, ra_orbit_sum, ra_sim, ra_sim_sum],\n", + " [dec_orbit, dec_orbit_sum, dec_sim, dec_sim_sum],\n", + "] = (times, rv_results, ra_results, dec_results)\n", "\n", "ra_gaia_err = np.full(np.shape(ra_sim_sum), sigma_ra_gaia)\n", "dec_gaia_err = np.full(np.shape(dec_sim_sum), sigma_dec_gaia)\n", "\n", "\n", - "\n", - "\n", - "\n", "# we assume that the first roman observation\n", "# will occur ~5 years after the last Gaia observation\n", - "t_1 = times_observed_astrometry_gaia[-1]+1800\n", + "t_1 = times_observed_astrometry_gaia[-1] + 1800\n", "\n", "# add roman observing times\n", "times_observed_astrometry_roman = []\n", - "for ii in range(t_1, t_1+(10*365)):\n", + "for ii in range(t_1, t_1 + (10 * 365)):\n", " if ii % 90 == 0:\n", - " times_observed_astrometry_roman.append(ii) \n", - "\n", + " times_observed_astrometry_roman.append(ii)\n", "\n", - "sigma_ra_roman = 5e-6 #5 micro-as\n", - "sigma_dec_roman = 5e-6 #5 micro-as\n", "\n", + "sigma_ra_roman = 5e-6 # 5 micro-as\n", + "sigma_dec_roman = 5e-6 # 5 micro-as\n", "\n", "\n", "times, rv_results, ra_results, dec_results = simulate_data(\n", - " n_planets, \n", - " sigma_rv, \n", + " n_planets,\n", + " sigma_rv,\n", " sigma_ra_roman,\n", " sigma_dec_roman,\n", " parallax,\n", " planet_params,\n", - " times_observed_rv = times_observed_rv,\n", - " times_observed_astrometry = times_observed_astrometry_roman\n", - " )\n", + " times_observed_rv=times_observed_rv,\n", + " times_observed_astrometry=times_observed_astrometry_roman,\n", + ")\n", "\n", "times_astrometry = np.append(times_astrometry, times[2], axis=0)\n", "\n", - "times_observed_astrometry = np.append(times_observed_astrometry, times[3], axis=0)\n", + "times_observed_astrometry = np.append(\n", + " times_observed_astrometry, times[3], axis=0\n", + ")\n", "\n", "ra_orbit = np.append(ra_orbit, ra_results[0], axis=0)\n", "ra_orbit_sum = np.append(ra_orbit_sum, ra_results[1], axis=0)\n", @@ -532,7 +537,7 @@ "dec_sim_sum = np.append(dec_sim_sum, dec_results[3], axis=0)\n", "\n", "ra_roman_err = np.full(np.shape(ra_results[3]), sigma_ra_roman)\n", - "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)\n" + "dec_roman_err = np.full(np.shape(dec_results[3]), sigma_dec_roman)" ] }, { @@ -552,36 +557,40 @@ "source": [ "import matplotlib.pyplot as plt\n", "import matplotlib\n", - "matplotlib.rc('font', **{'family': 'serif', 'serif': ['Computer Modern']})\n", - "matplotlib.rc('text', usetex=True)\n", + "\n", + "matplotlib.rc(\"font\", **{\"family\": \"serif\", \"serif\": [\"Computer Modern\"]})\n", + "matplotlib.rc(\"text\", usetex=True)\n", + "\n", "\n", "def plot_rv_signal(\n", - " n_planets, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv):\n", - " \n", - " \n", - " \n", - " fig, ax = plt.subplots(1, 1, figsize = [18,10], sharey='row')\n", - " fig.suptitle(\"RV Signal\", fontsize = 45)\n", + " n_planets,\n", + " rv_orbit,\n", + " rv_orbit_sum,\n", + " rv_sim,\n", + " rv_sim_sum,\n", + " times_rv,\n", + " times_observed_rv,\n", + "):\n", + " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", + " fig.suptitle(\"RV Signal\", fontsize=45)\n", " ax1 = ax\n", - " \n", - " \n", - " \n", "\n", - " ax1.plot(times_observed_rv, rv_sim_sum, 'o', color = 'k', label = 'combined signal', alpha = 0.3)\n", + " ax1.plot(\n", + " times_observed_rv,\n", + " rv_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " alpha=0.3,\n", + " )\n", "\n", - " ax1.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax1.set_ylabel(\"RV [m/s]\", fontsize = 27)\n", - " ax1.legend(fontsize = 27, loc = 2)\n", + " ax1.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax1.set_ylabel(\"RV [m/s]\", fontsize=27)\n", + " ax1.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -593,53 +602,55 @@ "outputs": [], "source": [ "def plot_astrometry_signal(\n", - " n_planets, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " n_planets,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry):\n", - " \n", - " \n", - "\n", - " fig, ax = plt.subplots(2, 1, figsize = [13,10], sharex='col', sharey='row')\n", - " fig.suptitle(\"Astrometric Signal\", fontsize = 45)\n", + " times_observed_astrometry,\n", + "):\n", + " fig, ax = plt.subplots(2, 1, figsize=[13, 10], sharex=\"col\", sharey=\"row\")\n", + " fig.suptitle(\"Astrometric Signal\", fontsize=45)\n", " ax2 = ax[0]\n", " ax3 = ax[1]\n", - " \n", - "\n", "\n", - " \n", - " ax2.plot(times_observed_astrometry, ra_sim_sum, 'o', color = 'k', label = 'combined signal')\n", + " ax2.plot(\n", + " times_observed_astrometry,\n", + " ra_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", " for tick in ax3.get_xticklabels():\n", " tick.set_rotation(30)\n", - " ax2.legend(fontsize = 27, loc = 2)\n", - " \n", - " ax3.plot(times_observed_astrometry, dec_sim_sum, 'o', color = 'k', label = 'combined signal')\n", - " ax2.set_xlabel(\"time [BJD]\", fontsize = 27)\n", - " ax3.legend(fontsize = 27, loc = 2)\n", - "\n", - "\n", + " ax2.legend(fontsize=27, loc=2)\n", + "\n", + " ax3.plot(\n", + " times_observed_astrometry,\n", + " dec_sim_sum,\n", + " \"o\",\n", + " color=\"k\",\n", + " label=\"combined signal\",\n", + " )\n", + " ax2.set_xlabel(\"time [BJD]\", fontsize=27)\n", + " ax3.legend(fontsize=27, loc=2)\n", "\n", " fig.tight_layout()\n", " fig.show()\n", "\n", - "\n", - "\n", " #######\n", " #######\n", " #######\n", "\n", - "\n", - " fig, ax = plt.subplots(1, 1, figsize = [10, 10])\n", + " fig, ax = plt.subplots(1, 1, figsize=[10, 10])\n", "\n", " # plot simulated data\n", - " ax.plot(ra_sim_sum, dec_sim_sum, 'o', color = 'k', label = \"combined signal\")\n", + " ax.plot(ra_sim_sum, dec_sim_sum, \"o\", color=\"k\", label=\"combined signal\")\n", "\n", " ax.set_ylabel(r\"$\\Delta \\delta$ ['']\")\n", " ax.set_xlabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\")\n", @@ -649,7 +660,7 @@ " ax.set_title(\"initial orbit\")\n", " ax.legend()\n", " fig.show()\n", - " \n", + "\n", " return None" ] }, @@ -665,46 +676,46 @@ "cell_type": "code", "execution_count": null, "id": "82e18e64", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_rv_signal(\n", - " 2, \n", - " rv_orbit, \n", - " rv_orbit_sum, \n", - " rv_sim, \n", - " rv_sim_sum, \n", - " times_rv, \n", - " times_observed_rv)\n", - " \n", - " \n" + " 2, rv_orbit, rv_orbit_sum, rv_sim, rv_sim_sum, times_rv, times_observed_rv\n", + ")" ] }, { "cell_type": "code", "execution_count": null, "id": "623b9c2d", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "plot_astrometry_signal(\n", - " 2, \n", - " ra_orbit, \n", - " ra_orbit_sum, \n", - " ra_sim, \n", - " ra_sim_sum, \n", - " dec_orbit, \n", - " dec_orbit_sum, \n", - " dec_sim, \n", + " 2,\n", + " ra_orbit,\n", + " ra_orbit_sum,\n", + " ra_sim,\n", + " ra_sim_sum,\n", + " dec_orbit,\n", + " dec_orbit_sum,\n", + " dec_sim,\n", " dec_sim_sum,\n", " times_astrometry,\n", - " times_observed_astrometry)" + " times_observed_astrometry,\n", + ")" ] }, { "cell_type": "markdown", "id": "99110cc0", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Ok, now we have a simulated dataset -- let's model it!\n", "\n", @@ -725,9 +736,7 @@ " print(\"------------\")\n", "\n", " with pm.Model() as model:\n", - " \n", - " \n", - " #log normal prior on period around estimates\n", + " # log normal prior on period around estimates\n", " logP = pm.Uniform(\n", " \"logP\",\n", " lower=0,\n", @@ -735,37 +744,33 @@ " shape=2,\n", " testval=np.log(periods),\n", " )\n", - " \n", - "\n", "\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", "\n", - "\n", " ## wide uniform prior on t_periastron\n", " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\n", - " \n", - " \n", + "\n", " # Wide normal prior for semi-amplitude\n", - " logK = pm.Uniform(\"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks))\n", - " \n", + " logK = pm.Uniform(\n", + " \"logK\", lower=-4, upper=4, shape=2, testval=np.log(Ks)\n", + " )\n", + "\n", " K = pm.Deterministic(\"K\", tt.exp(logK))\n", - " \n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), testval=0.01 * np.ones((2, 2)))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", "\n", - "\n", - "\n", - " \n", " # Jitter & a quadratic RV trend\n", - " #logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", + " # logs = pm.Normal(\"logs\", mu=np.log(np.median(y_rv_err)), sd=y_rv_err)\n", "\n", " # Then we define the orbit\n", - " orbit = xo.orbits.KeplerianOrbit(period=P, t_periastron=tperi, ecc=ecc, omega=omega)\n", + " orbit = xo.orbits.KeplerianOrbit(\n", + " period=P, t_periastron=tperi, ecc=ecc, omega=omega\n", + " )\n", "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", @@ -776,8 +781,6 @@ " # Sum over planets and add the background to get the full model\n", " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", "\n", - " \n", - "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", "\n", @@ -786,19 +789,17 @@ "\n", " # Finally add in the observation model. This next line adds a new contribution\n", " # to the log probability of the PyMC3 model\n", - " #err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", + " # err = tt.sqrt(y_rv_err ** 2 + tt.exp(2 * logs))\n", " err = y_rv_err\n", " pm.Normal(\"obs\", mu=rv_model, sd=err, observed=y_rv)\n", "\n", - "\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(start=map_soln, vars=[tperi])\n", " map_soln = pmx.optimize(start=map_soln, vars=[P])\n", " map_soln = pmx.optimize(start=map_soln, vars=[ecs])\n", " map_soln = pmx.optimize(start=map_soln)\n", "\n", - "\n", - " #return the max a-posteriori solution\n", + " # return the max a-posteriori solution\n", " return map_soln" ] }, @@ -806,16 +807,18 @@ "cell_type": "code", "execution_count": null, "id": "4d6faa40", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import pymc3 as pm\n", "import pymc3_ext as pmx\n", "\n", - "np.random.seed(979) \n", + "np.random.seed(979)\n", "\n", "\n", - "#rename variables in more consistent way for modeling\n", + "# rename variables in more consistent way for modeling\n", "x_rv = np.array(times_observed_rv)\n", "y_rv = rv_sim_sum\n", "y_rv_err = np.full(np.shape(y_rv), sigma_rv)\n", @@ -829,25 +832,29 @@ "dec_err = np.concatenate((dec_gaia_err, dec_roman_err))\n", "\n", "\n", - " \n", - " \n", "# make a fine grid that spans the observation window for plotting purposes\n", - "t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + "t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + ")\n", "t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", "# for predicted orbits\n", - "t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + "t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + ")\n", "\n", "\n", - "#minimize on RV data\n", + "# minimize on RV data\n", "periods_real = [4327.0, 365.0]\n", "periods_guess = []\n", "\n", - "#let's assume that THE can recover the period within 10% of the true value\n", + "# let's assume that THE can recover the period within 10% of the true value\n", "for per in periods_real:\n", - " periods_guess.append(per + np.random.uniform(low=-0.1*per, high=0.1*per))\n", + " periods_guess.append(\n", + " per + np.random.uniform(low=-0.1 * per, high=0.1 * per)\n", + " )\n", "\n", - "print('periods used as starting point = ', periods_guess)\n", + "print(\"periods used as starting point = \", periods_guess)\n", "periods_guess = np.array(periods_guess)\n", "\n", "\n", @@ -859,7 +866,9 @@ { "cell_type": "markdown", "id": "efc4dcc2", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## Now let's define out joint astrometry and radial velocity model:" ] @@ -871,25 +880,40 @@ "metadata": {}, "outputs": [], "source": [ - "def minimize_both(rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, ra_data, ra_err, dec_data, dec_err, parallax):\n", - " m_sun = 333030 #earth masses\n", - " \n", - " P_RV = np.array(rv_map_soln['P'])\n", - " K_RV = np.array(rv_map_soln['K'])\n", - " tperi_RV = np.array(rv_map_soln['tperi'])\n", - " ecc_RV = np.array(rv_map_soln['ecc'])\n", - " omega_RV = np.array(rv_map_soln['omega'])\n", - " min_masses_RV = xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value*317.83 #in m_earth\n", + "def minimize_both(\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", + "):\n", + " m_sun = 333030 # earth masses\n", + "\n", + " P_RV = np.array(rv_map_soln[\"P\"])\n", + " K_RV = np.array(rv_map_soln[\"K\"])\n", + " tperi_RV = np.array(rv_map_soln[\"tperi\"])\n", + " ecc_RV = np.array(rv_map_soln[\"ecc\"])\n", + " omega_RV = np.array(rv_map_soln[\"omega\"])\n", + " min_masses_RV = (\n", + " xo.estimate_minimum_mass(P_RV, x_rv, y_rv, y_rv_err).value * 317.83\n", + " ) # in m_earth\n", " phase_RV = determine_phase(P_RV, tperi_RV)\n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - "\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", " print(\"RV Solutions\")\n", " print(\"------------\")\n", @@ -899,98 +923,106 @@ " print(\"eccentricity: \", ecc_RV)\n", " print(\"omega: \", omega_RV)\n", "\n", - " print('')\n", + " print(\"\")\n", " print(\"minimizing joint model solutions pre-MCMC\")\n", " print(\"------------\")\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", - " inc_test_vals = np.array(np.radians([5., 25., 45., 65., 85.]))\n", + " inc_test_vals = np.array(np.radians([5.0, 25.0, 45.0, 65.0, 85.0]))\n", " model, map_soln, logp = [], [], []\n", " for inc in inc_test_vals:\n", - " mass_test_vals = min_masses_RV/np.sin(inc)\n", - " print('')\n", + " mass_test_vals = min_masses_RV / np.sin(inc)\n", + " print(\"\")\n", " print(\"trying inclination = \" + str(np.degrees(inc)))\n", " print(\"mass test val = \" + str(mass_test_vals))\n", " print(\"------------\")\n", - " \n", - " \n", - "\n", - "\n", "\n", " def get_model():\n", " with pm.Model() as model:\n", - "\n", - "\n", " # Below we will run a version of this model where a measurement of parallax is provided\n", " # The measurement is in milliarcsec\n", " m_plx = pm.Bound(pm.Normal, lower=0, upper=200)(\n", - " \"m_plx\", mu=parallax*1000, sd=10, testval=parallax*1000\n", + " \"m_plx\", mu=parallax * 1000, sd=10, testval=parallax * 1000\n", " )\n", " plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n", "\n", - "\n", " # We expect the period to be around that found from just the RVs,\n", " # so we'll set a broad prior on logP\n", - " \n", + "\n", " logP = pm.Uniform(\n", - " \"logP\", lower=0, upper=np.log(2*P_RV), testval=np.log(P_RV), shape=2\n", + " \"logP\",\n", + " lower=0,\n", + " upper=np.log(2 * P_RV),\n", + " testval=np.log(P_RV),\n", + " shape=2,\n", " )\n", " P = pm.Deterministic(\"P\", tt.exp(logP))\n", - " \n", + "\n", " # Eccentricity & argument of periasteron\n", - " ecs = pmx.UnitDisk(\"ecs\", shape=(2, 2), \n", - " testval=np.array([np.sqrt(ecc_RV)*np.cos(omega_RV), \n", - " np.sqrt(ecc_RV)*np.sin(omega_RV)]))\n", - " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs ** 2, axis=0))\n", + " ecs = pmx.UnitDisk(\n", + " \"ecs\",\n", + " shape=(2, 2),\n", + " testval=np.array(\n", + " [\n", + " np.sqrt(ecc_RV) * np.cos(omega_RV),\n", + " np.sqrt(ecc_RV) * np.sin(omega_RV),\n", + " ]\n", + " ),\n", + " )\n", + " ecc = pm.Deterministic(\"ecc\", tt.sum(ecs**2, axis=0))\n", " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\n", "\n", - " xo.eccentricity.kipping13(\"ecc_prior\", fixed=True, observed=ecc)\n", - "\n", - " \n", - " \n", + " xo.eccentricity.kipping13(\n", + " \"ecc_prior\", fixed=True, observed=ecc\n", + " )\n", "\n", " # Omegas are co-dependent, so sample them with variables Omega_plus\n", - " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and \n", + " # and Omegas_minus. Omega_plus is (Omega_0 + Omega_1)/2 and\n", " # Omega_minus is (Omega_0 - Omega_1)/2\n", - " \n", + "\n", " Omega_plus = pmx.Angle(\"Omega_plus\", shape=1)\n", " Omega_minus = pmx.Angle(\"Omega_minus\", shape=1)\n", - " \n", - " \n", - " Omega = tt.concatenate( [(Omega_plus + Omega_minus),\n", - " (Omega_plus - Omega_minus)] )\n", - " \n", "\n", - " Omega = pm.Deterministic(\"Omega\", Omega) \n", - " Omega_sum = pm.Deterministic(\"Omega_sum\", ((Omega_plus)*2)% np.pi)\n", - " Omega_diff = pm.Deterministic(\"Omega_diff\", ((Omega_minus)*2)% np.pi)\n", - " \n", + " Omega = tt.concatenate(\n", + " [(Omega_plus + Omega_minus), (Omega_plus - Omega_minus)]\n", + " )\n", "\n", + " Omega = pm.Deterministic(\"Omega\", Omega)\n", + " Omega_sum = pm.Deterministic(\n", + " \"Omega_sum\", ((Omega_plus) * 2) % np.pi\n", + " )\n", + " Omega_diff = pm.Deterministic(\n", + " \"Omega_diff\", ((Omega_minus) * 2) % np.pi\n", + " )\n", "\n", - " \n", " # For these orbits, it can also be better to fit for a phase angle\n", " # (relative to a reference time) instead of the time of periasteron\n", " # passage directly\n", " phase = pmx.Angle(\"phase\", testval=phase_RV, shape=2)\n", " tperi = pm.Deterministic(\"tperi\", P * phase / (2 * np.pi))\n", - " \n", "\n", - " \n", " # uniform prior on sqrtm_sini and sqrtm_cosi (upper 100* testval to stop planet flipping)\n", - " log_m = pm.Uniform(\"log_m\", lower=-1, upper=np.log(100*mass_test_vals), testval=np.log(mass_test_vals), shape=2)\n", + " log_m = pm.Uniform(\n", + " \"log_m\",\n", + " lower=-1,\n", + " upper=np.log(100 * mass_test_vals),\n", + " testval=np.log(mass_test_vals),\n", + " shape=2,\n", + " )\n", " m_planet = pm.Deterministic(\"m_planet\", tt.exp(log_m))\n", - " m_planet_fit = pm.Deterministic(\"m_planet_fit\", m_planet/m_sun)\n", - "\n", + " m_planet_fit = pm.Deterministic(\n", + " \"m_planet_fit\", m_planet / m_sun\n", + " )\n", "\n", - " cos_incl = pm.Uniform(\"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2)\n", + " cos_incl = pm.Uniform(\n", + " \"cos_incl\", lower=0, upper=1, testval=np.cos(inc), shape=2\n", + " )\n", " incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n", - " \n", - "\n", "\n", - " \n", - " \n", " # Set up the orbit\n", " orbit = xo.orbits.KeplerianOrbit(\n", " t_periastron=tperi,\n", @@ -999,13 +1031,10 @@ " ecc=ecc,\n", " omega=omega,\n", " Omega=Omega,\n", - " m_planet = m_planet_fit,\n", - " plx=plx\n", + " m_planet=m_planet_fit,\n", + " plx=plx,\n", " )\n", "\n", - "\n", - " \n", - " \n", " # Add a function for computing the full astrometry model\n", " def get_astrometry_model(t, name=\"\"):\n", " # First the astrometry induced by the planets\n", @@ -1013,54 +1042,56 @@ " # determine and print the star position at desired times\n", " pos = orbit.get_star_position(t, plx)\n", "\n", - " x,y,z = pos\n", - "\n", + " x, y, z = pos\n", "\n", " # calculate rho and theta\n", - " rhos = tt.squeeze(tt.sqrt(x ** 2 + y ** 2)) # arcsec\n", - " thetas = tt.squeeze(tt.arctan2(y, x)) # radians between [-pi, pi]\n", - " \n", - " \n", - " #rhos, thetas = get_star_relative_angles(t, plx)\n", - " \n", - " \n", - " dec = pm.Deterministic(\"dec\" + name, rhos * np.cos(thetas)) # X is north\n", - " ra = pm.Deterministic(\"ra\" + name, rhos * np.sin(thetas)) # Y is east\n", - " \n", + " rhos = tt.squeeze(tt.sqrt(x**2 + y**2)) # arcsec\n", + " thetas = tt.squeeze(\n", + " tt.arctan2(y, x)\n", + " ) # radians between [-pi, pi]\n", + "\n", + " # rhos, thetas = get_star_relative_angles(t, plx)\n", + "\n", + " dec = pm.Deterministic(\n", + " \"dec\" + name, rhos * np.cos(thetas)\n", + " ) # X is north\n", + " ra = pm.Deterministic(\n", + " \"ra\" + name, rhos * np.sin(thetas)\n", + " ) # Y is east\n", + "\n", " # Sum over planets to get the full model\n", - " dec_model = pm.Deterministic(\"dec_model\" + name, tt.sum(dec, axis=-1))\n", - " ra_model = pm.Deterministic(\"ra_model\" + name, tt.sum(ra, axis=-1))\n", - " \n", + " dec_model = pm.Deterministic(\n", + " \"dec_model\" + name, tt.sum(dec, axis=-1)\n", + " )\n", + " ra_model = pm.Deterministic(\n", + " \"ra_model\" + name, tt.sum(ra, axis=-1)\n", + " )\n", "\n", - " \n", " return dec_model, ra_model\n", "\n", - " \n", " # Define the astrometry model at the observed times\n", " dec_model, ra_model = get_astrometry_model(x_astrometry)\n", "\n", " # Also define the model on a fine grid as computed above (for plotting)\n", - " dec_model_fine, ra_model_fine = get_astrometry_model(t_fine, name=\"_fine\")\n", - "\n", - " \n", - "\n", + " dec_model_fine, ra_model_fine = get_astrometry_model(\n", + " t_fine, name=\"_fine\"\n", + " )\n", "\n", - " #dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", - " #ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", + " # dec_tot_err = tt.sqrt(dec_err ** 2 + tt.exp(2 * log_dec_s))\n", + " # ra_tot_err = tt.sqrt(ra_err ** 2 + tt.exp(2 * log_ra_s))\n", " dec_tot_err = dec_err\n", " ra_tot_err = ra_err\n", "\n", + " # define the likelihood function, e.g., a Gaussian on both ra and dec\n", + " pm.Normal(\n", + " \"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data\n", + " )\n", + " pm.Normal(\n", + " \"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data\n", + " )\n", "\n", - " # define the likelihood function, e.g., a Gaussian on both ra and dec \n", - " pm.Normal(\"dec_obs\", mu=dec_model, sd=dec_tot_err, observed=dec_data)\n", - " pm.Normal(\"ra_obs\", mu=ra_model, sd=ra_tot_err, observed=ra_data)\n", - "\n", - "\n", - " \n", - " \n", " # ADD RV MODEL\n", "\n", - "\n", " # And a function for computing the full RV model\n", " def get_rv_model(t, name=\"\"):\n", " # First the RVs induced by the planets\n", @@ -1068,7 +1099,9 @@ " pm.Deterministic(\"vrad\" + name, vrad)\n", "\n", " # Sum over planets to get the full model\n", - " return pm.Deterministic(\"rv_model\" + name, tt.sum(vrad, axis=-1))\n", + " return pm.Deterministic(\n", + " \"rv_model\" + name, tt.sum(vrad, axis=-1)\n", + " )\n", "\n", " # Define the RVs at the observed times\n", " rv_model = get_rv_model(x_rv)\n", @@ -1084,24 +1117,20 @@ " # Optimize to find the initial parameters\n", " map_soln = model.test_point\n", " map_soln = pmx.optimize(map_soln, vars=[Omega, phase])\n", - " #map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", + " # map_soln = pmx.optimize(map_soln, vars=[Omega, m_planet, incl, ecs, phase])\n", "\n", " map_soln = pmx.optimize(map_soln)\n", "\n", - "\n", " return model, map_soln\n", "\n", - "\n", - "\n", " a_model, a_map_soln = get_model()\n", - " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis = 0)\n", - " print('log likelihood = ' + str(a_logp))\n", + " a_logp = a_model.check_test_point(test_point=a_map_soln).sum(axis=0)\n", + " print(\"log likelihood = \" + str(a_logp))\n", "\n", " model.append(a_model)\n", " map_soln.append(a_map_soln)\n", " logp.append(a_logp)\n", "\n", - " \n", " best_index = 0\n", " for index in range(0, len(model)):\n", " if logp[index] >= logp[best_index]:\n", @@ -1111,31 +1140,26 @@ " the_map_soln = map_soln[best_index]\n", " the_logp = logp[best_index]\n", "\n", - "\n", - "\n", " return the_model, the_map_soln, the_logp\n", "\n", "\n", - "\n", "def determine_phase(P, t_periastron):\n", " phase = (2 * np.pi * t_periastron) / P\n", " return phase\n", "\n", "\n", - "\n", "def model_both(model, map_soln, tune_steps, draw_steps):\n", - " print('Joint RV + Astometry Minimization Solutions:')\n", + " print(\"Joint RV + Astometry Minimization Solutions:\")\n", " print(\"------------\")\n", "\n", - " print('m_planet: ' + str(map_soln['m_planet']))\n", - " print('P: ' + str(map_soln['P']))\n", - " print('incl: ' + str(map_soln['incl']))\n", - " print('Omega: ' + str(map_soln['Omega']))\n", - " print('tperi: ' + str(map_soln['tperi']))\n", - " print('ecc: ' + str(map_soln['ecc']))\n", - " print('omega: ' + str(map_soln['omega']))\n", - " print('plx: ' + str(map_soln['plx']))\n", - "\n", + " print(\"m_planet: \" + str(map_soln[\"m_planet\"]))\n", + " print(\"P: \" + str(map_soln[\"P\"]))\n", + " print(\"incl: \" + str(map_soln[\"incl\"]))\n", + " print(\"Omega: \" + str(map_soln[\"Omega\"]))\n", + " print(\"tperi: \" + str(map_soln[\"tperi\"]))\n", + " print(\"ecc: \" + str(map_soln[\"ecc\"]))\n", + " print(\"omega: \" + str(map_soln[\"omega\"]))\n", + " print(\"plx: \" + str(map_soln[\"plx\"]))\n", "\n", " with model:\n", " trace = pmx.sample(\n", @@ -1163,18 +1187,24 @@ "\n", "################\n", "################\n", - "#minimize on joint model\n", + "# minimize on joint model\n", "joint_model, joint_map_soln, joint_logp = minimize_both(\n", - " rv_map_soln, x_rv, y_rv, y_rv_err, x_astrometry, \n", - " ra_data, ra_err, dec_data, dec_err, parallax\n", + " rv_map_soln,\n", + " x_rv,\n", + " y_rv,\n", + " y_rv_err,\n", + " x_astrometry,\n", + " ra_data,\n", + " ra_err,\n", + " dec_data,\n", + " dec_err,\n", + " parallax,\n", ")\n", "\n", "\n", - "\n", - "\n", "################\n", "################\n", - "#run full MCMC\n", + "# run full MCMC\n", "trace = model_both(joint_model, joint_map_soln, 1000, 1000)" ] }, @@ -1193,143 +1223,168 @@ "metadata": {}, "outputs": [], "source": [ - "matplotlib.rc('text', usetex=False)\n", + "matplotlib.rc(\"text\", usetex=False)\n", + "\n", "\n", "def make_plots(trace, orbit_params, n_planets):\n", - " m_sun = 333030 #earth masses\n", + " m_sun = 333030 # earth masses\n", "\n", " [orbit_params_earth, orbit_params_jup] = orbit_params\n", "\n", - "\n", - " [P_earth, e_earth, Tper_earth, omega_earth, Omega_earth, \n", - " inclination_earth, m_earth] = orbit_params_earth\n", - "\n", - "\n", - " [P_jup, e_jup, Tper_jup, omega_jup, Omega_jup, \n", - " inclination_jup, m_jup] = orbit_params_jup\n", - "\n", - "\n", - " a_true_earth = a_from_Kepler3(P_earth, 1.0+m_earth)\n", - " a_true_jup = a_from_Kepler3(P_jup, 1.0+m_jup)\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", - " truths = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth, \n", - " Omega_earth-Omega_jup, Omega_earth+Omega_jup, inclination_jup, inclination_earth, \n", - " 0.1, m_jup*m_sun, m_earth*m_sun, Tper_jup, Tper_earth]\n", - "\n", - "\n", - " #[P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", - " truth_chain_plot = [P_jup, P_earth, e_jup, e_earth, omega_jup, omega_earth,\n", - " inclination_jup, inclination_earth, Tper_jup, Tper_earth, m_jup*m_sun, m_earth*m_sun]\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " [\n", + " P_earth,\n", + " e_earth,\n", + " Tper_earth,\n", + " omega_earth,\n", + " Omega_earth,\n", + " inclination_earth,\n", + " m_earth,\n", + " ] = orbit_params_earth\n", + "\n", + " [\n", + " P_jup,\n", + " e_jup,\n", + " Tper_jup,\n", + " omega_jup,\n", + " Omega_jup,\n", + " inclination_jup,\n", + " m_jup,\n", + " ] = orbit_params_jup\n", + "\n", + " a_true_earth = a_from_Kepler3(P_earth, 1.0 + m_earth)\n", + " a_true_jup = a_from_Kepler3(P_jup, 1.0 + m_jup)\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, Omega_sum, Omega_minus, incl1, incl2, plx, m1, m2, a1, a2, Tper1, Tper2]\n", + " truths = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " Omega_earth - Omega_jup,\n", + " Omega_earth + Omega_jup,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " 0.1,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " ]\n", + "\n", + " # [P1, P2, e1, e2, omega1, omega2, incl1, incl2, Tper1, Tper2, m1, m2, a1, a2]\n", + " truth_chain_plot = [\n", + " P_jup,\n", + " P_earth,\n", + " e_jup,\n", + " e_earth,\n", + " omega_jup,\n", + " omega_earth,\n", + " inclination_jup,\n", + " inclination_earth,\n", + " Tper_jup,\n", + " Tper_earth,\n", + " m_jup * m_sun,\n", + " m_earth * m_sun,\n", + " ]\n", "\n", " # plot the table summarizing MCMC results\n", " az.summary(\n", " trace,\n", - " var_names=[\"P\", \"tperi\", \"omega\", \"Omega_sum\", \"Omega_plus\", \"Omega_minus\", \n", - " \"incl\", \"ecc\", \"plx\", \"m_planet\"],\n", + " var_names=[\n", + " \"P\",\n", + " \"tperi\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_plus\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"ecc\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " ],\n", " )\n", "\n", - "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", " # plot the corner plots\n", " _ = corner.corner(\n", - " trace, var_names=[\"P\", \"ecc\", \"omega\", \"Omega_sum\", \"Omega_minus\", \"incl\", \n", - " \"plx\", \"m_planet\", 'tperi'], quantiles=[0.16, 0.5, 0.84],\n", - " show_titles=True, title_kwargs={\"fontsize\": 13}, \n", - " truths = truths, truth_color = \"#03003a\"\n", + " trace,\n", + " var_names=[\n", + " \"P\",\n", + " \"ecc\",\n", + " \"omega\",\n", + " \"Omega_sum\",\n", + " \"Omega_minus\",\n", + " \"incl\",\n", + " \"plx\",\n", + " \"m_planet\",\n", + " \"tperi\",\n", + " ],\n", + " quantiles=[0.16, 0.5, 0.84],\n", + " show_titles=True,\n", + " title_kwargs={\"fontsize\": 13},\n", + " truths=truths,\n", + " truth_color=\"#03003a\",\n", " )\n", "\n", " plt.show()\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the chains\n", " parameters = [\"P\", \"ecc\", \"omega\", \"incl\", \"tperi\", \"m_planet\"]\n", " for ii in range(0, len(parameters)):\n", " plot_truth = False\n", " param = parameters[ii]\n", - " \n", - " true_vals_earth = truth_chain_plot[2*ii]\n", - " true_vals_jup = truth_chain_plot[2*ii+1]\n", + "\n", + " true_vals_earth = truth_chain_plot[2 * ii]\n", + " true_vals_jup = truth_chain_plot[2 * ii + 1]\n", " plot_truth = True\n", - " \n", - " fig, ax = plt.subplots(1,2, figsize = (15,3))\n", + "\n", + " fig, ax = plt.subplots(1, 2, figsize=(15, 3))\n", " planet1_chain1 = trace.posterior[param].values[:, :, 0][0]\n", " planet1_chain2 = trace.posterior[param].values[:, :, 0][1]\n", - " \n", + "\n", " planet2_chain1 = trace.posterior[param].values[:, :, 1][0]\n", " planet2_chain2 = trace.posterior[param].values[:, :, 1][1]\n", - " \n", - " \n", - " nstep = np.arange(1, len(planet1_chain1)+1, 1)\n", - " \n", - " \n", + "\n", + " nstep = np.arange(1, len(planet1_chain1) + 1, 1)\n", + "\n", " ax[0].plot(nstep, planet1_chain1)\n", " ax[0].plot(nstep, planet1_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[0].axhline(y=true_vals_earth, color = 'r', label = 'truth')\n", - " ax[0].set_title(\"Jupiter\", fontsize = 18)\n", - " ax[0].legend(fontsize = 18)\n", - " \n", + " ax[0].axhline(y=true_vals_earth, color=\"r\", label=\"truth\")\n", + " ax[0].set_title(\"Jupiter\", fontsize=18)\n", + " ax[0].legend(fontsize=18)\n", + "\n", " ax[1].plot(nstep, planet2_chain1)\n", " ax[1].plot(nstep, planet2_chain2)\n", - " \n", + "\n", " if plot_truth:\n", - " ax[1].axhline(y=true_vals_jup, color = 'r', label = 'truth')\n", - " ax[1].set_title(\"Earth\", fontsize = 18)\n", - " ax[1].legend(fontsize = 18)\n", + " ax[1].axhline(y=true_vals_jup, color=\"r\", label=\"truth\")\n", + " ax[1].set_title(\"Earth\", fontsize=18)\n", + " ax[1].legend(fontsize=18)\n", "\n", - " fig.suptitle(param, fontsize = 18)\n", + " fig.suptitle(param, fontsize=18)\n", " fig.tight_layout()\n", " plt.show()\n", "\n", + " return \"plotting complete\"\n", "\n", "\n", + "def a_from_Kepler3(period, M_tot):\n", + " period = period * 86400 # days to seconds\n", "\n", - " return \"plotting complete\"\n", - "\n", + " M_tot = M_tot * M_sun.value # solar masses to kg\n", "\n", + " a3 = (((constants.G.value) * M_tot) / (4 * np.pi**2)) * (period) ** 2.0\n", "\n", + " a = a3 ** (1 / 3)\n", "\n", + " a = a * 6.68459 * 10 ** (-12.0) # meters to AU\n", "\n", - "def a_from_Kepler3(period, M_tot):\n", - " period = period*86400 #days to seconds\n", - " \n", - "\n", - " M_tot = M_tot*M_sun.value #solar masses to kg\n", - " \n", - " a3 = ( ((constants.G.value)*M_tot) / (4*np.pi**2) ) * (period)**2.\n", - " \n", - " a = a3**(1/3)\n", - " \n", - " \n", - " \n", - " a = a * 6.68459*10**(-12.) # meters to AU\n", - " \n", - " return(a) #in AUs\n", - "\n", - "\n", - "\n" + " return a # in AUs" ] }, { @@ -1344,18 +1399,14 @@ "cell_type": "code", "execution_count": null, "id": "cfd9e0cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import arviz as az\n", "\n", - "az.summary(trace, var_names=[\n", - " \"P\",\n", - " \"ecc\",\n", - " \"m_planet\",\n", - " \"Omega\",\n", - " \"omega\",\n", - " \"incl\"])\n" + "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" ] }, { @@ -1372,12 +1423,14 @@ "cell_type": "code", "execution_count": null, "id": "6775ae19", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ "import corner\n", "\n", - "make_plots(trace, planet_params, 2)\n" + "make_plots(trace, planet_params, 2)" ] }, { @@ -1396,34 +1449,31 @@ "outputs": [], "source": [ "import pandas as pd\n", - "#save simulated data as dataframe \n", + "\n", + "# save simulated data as dataframe\n", "simulated_data_dic = {\n", " \"times_rv_observed\": x_rv,\n", " \"rv_observed\": y_rv,\n", " \"rv_err_observed\": y_rv_err,\n", - "\n", " \"times_rv_orbit\": times_rv,\n", " \"rv_orbit\": rv_orbit,\n", " \"rv_orbit_sum\": rv_orbit_sum,\n", - "\n", " \"times_astrometry_observed\": x_astrometry,\n", " \"ra_observed\": ra_data,\n", " \"dec_observed\": dec_data,\n", " \"ra_err_observed\": ra_err,\n", " \"dec_err_observed\": dec_err,\n", - "\n", " \"times_astrometry_orbit\": times_astrometry,\n", " \"ra_orbit\": ra_orbit,\n", " \"ra_orbit_sum\": ra_orbit_sum,\n", " \"dec_orbit\": dec_orbit,\n", - " \"dec_orbit_sum\": dec_orbit_sum\n", - " }\n", + " \"dec_orbit_sum\": dec_orbit_sum,\n", + "}\n", "\n", "\n", - "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient='index')\n", + "simulated_data = pd.DataFrame.from_dict(simulated_data_dic, orient=\"index\")\n", "simulated_data = simulated_data.transpose()\n", - "simulated_data = simulated_data.fillna(value=np.nan)\n", - "\n" + "simulated_data = simulated_data.fillna(value=np.nan)" ] }, { @@ -1435,55 +1485,48 @@ "source": [ "def remove_nans(array):\n", " array = array[~pd.isnull(array)]\n", - " \n", - " return array\n", - "\n", "\n", + " return array\n", "\n", "\n", "def import_orbit_data(orbit_row):\n", " import math\n", - " \n", + "\n", " orbit_row = remove_nans(orbit_row)\n", " orbit_array = []\n", " for row in orbit_row:\n", " orbit_array.append([np.array(row[1]), row[0]])\n", "\n", " orbit_array = np.array(orbit_array)\n", - " \n", - " return orbit_array\n", - "\n", "\n", + " return orbit_array\n", "\n", "\n", "def make_rv_astrometry_plots(simulated_data, trace, nplanets):\n", " import matplotlib.pyplot as plt\n", "\n", - " \n", - " \n", - " #simulated RV observations\n", - " x_rv = simulated_data['times_rv_observed'].values\n", - " y_rv = simulated_data['rv_observed'].values\n", - " y_rv_err = simulated_data['rv_err_observed'].values\n", - " \n", - " #simulated RV full orbits\n", - " x_rv_orbit = simulated_data['times_rv_orbit'].values\n", - " y_rv_orbit = simulated_data['rv_orbit'].values\n", - " \n", - " #simulated astrometric observations\n", - " x_astrometry = simulated_data['times_astrometry_observed'].values\n", - " y_ra = simulated_data['ra_observed'].values\n", - " y_dec = simulated_data['dec_observed'].values\n", - " y_ra_err = simulated_data['ra_err_observed'].values\n", - " y_dec_err = simulated_data['dec_err_observed'].values\n", - " \n", - " #simulated astrometric full orbits\n", - " x_astrometry_orbit = simulated_data['times_astrometry_orbit'].values\n", - " y_ra_orbit = simulated_data['ra_orbit'].values\n", - " y_dec_orbit = simulated_data['dec_orbit'].values\n", - " \n", - " \n", - " #remove any nans from simulated data (created because of mismatching lengths)\n", + " # simulated RV observations\n", + " x_rv = simulated_data[\"times_rv_observed\"].values\n", + " y_rv = simulated_data[\"rv_observed\"].values\n", + " y_rv_err = simulated_data[\"rv_err_observed\"].values\n", + "\n", + " # simulated RV full orbits\n", + " x_rv_orbit = simulated_data[\"times_rv_orbit\"].values\n", + " y_rv_orbit = simulated_data[\"rv_orbit\"].values\n", + "\n", + " # simulated astrometric observations\n", + " x_astrometry = simulated_data[\"times_astrometry_observed\"].values\n", + " y_ra = simulated_data[\"ra_observed\"].values\n", + " y_dec = simulated_data[\"dec_observed\"].values\n", + " y_ra_err = simulated_data[\"ra_err_observed\"].values\n", + " y_dec_err = simulated_data[\"dec_err_observed\"].values\n", + "\n", + " # simulated astrometric full orbits\n", + " x_astrometry_orbit = simulated_data[\"times_astrometry_orbit\"].values\n", + " y_ra_orbit = simulated_data[\"ra_orbit\"].values\n", + " y_dec_orbit = simulated_data[\"dec_orbit\"].values\n", + "\n", + " # remove any nans from simulated data (created because of mismatching lengths)\n", " x_rv = remove_nans(x_rv)\n", " y_rv = remove_nans(y_rv)\n", " y_rv_err = remove_nans(y_rv_err)\n", @@ -1497,43 +1540,34 @@ " x_astrometry_orbit = remove_nans(x_astrometry_orbit)\n", " y_ra_orbit = import_orbit_data(y_ra_orbit)\n", " y_dec_orbit = import_orbit_data(y_dec_orbit)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " # make a fine grid that spans the observation window for plotting purposes\n", - " t_astrometry = np.linspace(x_astrometry.min() - 5, x_astrometry.max() + 5, 1000)\n", + " t_astrometry = np.linspace(\n", + " x_astrometry.min() - 5, x_astrometry.max() + 5, 1000\n", + " )\n", " t_rv = np.linspace(x_rv.min() - 5, x_rv.max() + 5, 1000)\n", "\n", " # for predicted orbits\n", - " t_fine = np.linspace(x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000)\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " # add NaN for plotting purpose on true model\n", - " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float('NaN'))\n", - " y_ra_orbit = np.insert(y_ra_orbit, 10000, float('NaN'), axis=0)\n", - " y_dec_orbit = np.insert(y_dec_orbit, 10000, float('NaN'),axis=0)\n", - " \n", - " \n", - " \n", + " t_fine = np.linspace(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", "\n", + " # add NaN for plotting purpose on true model\n", + " x_astrometry_orbit = np.insert(x_astrometry_orbit, 10000, float(\"NaN\"))\n", + " y_ra_orbit = np.insert(y_ra_orbit, 10000, float(\"NaN\"), axis=0)\n", + " y_dec_orbit = np.insert(y_dec_orbit, 10000, float(\"NaN\"), axis=0)\n", "\n", - " \n", - " \n", " figsAs = []\n", " figsRV = []\n", " for n in range(0, nplanets):\n", - " \n", " ######\n", - " #make rv plots\n", - " \n", - " \n", - " figRV, ax = plt.subplots(2, 1, figsize = [13, 6], gridspec_kw={'height_ratios': [3, 1]})\n", + " # make rv plots\n", + "\n", + " figRV, ax = plt.subplots(\n", + " 2, 1, figsize=[13, 6], gridspec_kw={\"height_ratios\": [3, 1]}\n", + " )\n", " ax0, ax1 = ax[0], ax[1]\n", - " \n", + "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1544,86 +1578,68 @@ " trace.posterior[\"vrad\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # Plot the data\n", - " ax0.errorbar(x_rv, y_rv - other, yerr=y_rv_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax0.errorbar(\n", + " x_rv,\n", + " y_rv - other,\n", + " yerr=y_rv_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"vrad_pred\"].values[:, :, :, n],\n", " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_rv = []\n", " for an_x in x_rv:\n", " med_rv.append(np.interp(an_x, t_rv, pred[1]))\n", "\n", - "\n", - "\n", - " \n", " # plot the MCMC model\n", " ax0.plot(t_rv, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_rv, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_rv, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", - "\n", "\n", + " # plot the true RA model\n", + " ax0.plot(x_rv_orbit, y_rv_orbit.T[n], color=\"k\", label=\"truth\")\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"phase [days]\", fontsize=18)\n", " ax0.set_ylabel(\"radial velocity [m/s]\", fontsize=18)\n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", - " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - "\n", "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " \n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_rv, y_rv - other - med_rv,\n", - " yerr=y_rv_err, color = '#00257c', fmt=\".\"\n", + " x_rv,\n", + " y_rv - other - med_rv,\n", + " yerr=y_rv_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", - " \n", + "\n", " figRV.tight_layout()\n", " figRV.subplots_adjust(hspace=0.5)\n", - " figRV.suptitle(\"planet {0}\".format(n+1), fontsize=33)\n", + " figRV.suptitle(\"planet {0}\".format(n + 1), fontsize=33)\n", " figRV.tight_layout()\n", - " \n", - " \n", - "\n", - " \n", - " \n", - " \n", - " \n", - " \n", + "\n", " figRV.show()\n", " ################\n", - " #plot astrometry\n", - " \n", - " \n", - "\n", - " figAs, ax = plt.subplots(2, 2, figsize = [18,9], sharey='row', gridspec_kw={'height_ratios': [3, 1]})\n", + " # plot astrometry\n", + "\n", + " figAs, ax = plt.subplots(\n", + " 2,\n", + " 2,\n", + " figsize=[18, 9],\n", + " sharey=\"row\",\n", + " gridspec_kw={\"height_ratios\": [3, 1]},\n", + " )\n", " ax0, ax1, ax2, ax3 = ax[0][0], ax[1][0], ax[0][1], ax[1][1]\n", "\n", " # Get the posterior median orbital parameters\n", @@ -1636,14 +1652,16 @@ " trace.posterior[\"dec\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", "\n", - "\n", - "\n", - "\n", - "\n", - "\n", " # plot the DEC data\n", - " ax0.errorbar(x_astrometry, y_dec - other, yerr=y_dec_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", + " ax0.errorbar(\n", + " x_astrometry,\n", + " y_dec - other,\n", + " yerr=y_dec_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " # plot the MCMC DEC model\n", " pred = np.percentile(\n", @@ -1651,55 +1669,78 @@ " [16, 50, 84],\n", " axis=(0, 1),\n", " )\n", - " \n", + "\n", " med_dec = []\n", " for an_x in x_astrometry:\n", " med_dec.append(np.interp(an_x, t_fine, pred[1]))\n", - " \n", - "\n", "\n", " ax0.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax0.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax0.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", "\n", " # plot the true DEC model\n", - " ax0.plot(x_astrometry_orbit, y_dec_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - " \n", - " \n", - " ax0.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax0.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax0.transAxes)\n", - " ax0.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax0.transAxes)\n", - " ax0.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax0.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax0.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - "\n", - "\n", + " ax0.plot(\n", + " x_astrometry_orbit, y_dec_orbit.T[n], color=\"k\", label=\"truth\"\n", + " )\n", "\n", + " ax0.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax0.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax0.transAxes\n", + " )\n", + " ax0.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax0.transAxes\n", + " )\n", + " ax0.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax0.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax0.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", " ax0.legend(fontsize=9, loc=1)\n", " ax0.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize = 18)\n", - " \n", - " \n", - " ax1.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax0.set_ylabel(r\"$\\Delta \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax1.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax1.set_xlabel(\"phase [days]\", fontsize=18)\n", - " \n", - " \n", - " #plot residuals \n", + "\n", + " # plot residuals\n", " ax1.axhline(0.0, color=\"k\")\n", " ax1.errorbar(\n", - " x_astrometry, y_dec - other - med_dec,\n", - " yerr=y_dec_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_dec - other - med_dec,\n", + " yerr=y_dec_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", "\n", - "\n", - "\n", " # Get the posterior median orbital parameters\n", " p = np.median(trace.posterior[\"P\"].values[:, :, n])\n", " t0 = np.median(trace.posterior[\"tperi\"].values[:, :, n])\n", @@ -1709,86 +1750,99 @@ " other = np.median(\n", " trace.posterior[\"ra\"].values[:, :, :, (n + 1) % 2], axis=(0, 1)\n", " )\n", - " \n", - "\n", - "\n", - "\n", - "\n", "\n", " # Plot the RA data\n", - " ax2.errorbar(x_astrometry, y_ra - other, yerr=y_ra_err, fmt=\".\", \n", - " color = \"#00257c\", label = \"data\", alpha = 0.7)\n", - "\n", + " ax2.errorbar(\n", + " x_astrometry,\n", + " y_ra - other,\n", + " yerr=y_ra_err,\n", + " fmt=\".\",\n", + " color=\"#00257c\",\n", + " label=\"data\",\n", + " alpha=0.7,\n", + " )\n", "\n", " pred = np.percentile(\n", " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", " [16, 50, 84],\n", - " axis=(0, 1), \n", + " axis=(0, 1),\n", " )\n", - " \n", - " \n", + "\n", " med_ra = []\n", " for an_x in x_astrometry:\n", " med_ra.append(np.interp(an_x, t_fine, pred[1]))\n", "\n", - "\n", " # plot the MCMC RA model\n", " ax2.plot(t_fine, pred[1], color=\"r\", label=\"posterior\")\n", - " art = ax2.fill_between(\n", - " t_fine, pred[0], pred[2], color=\"r\", alpha=0.3\n", - " )\n", + " art = ax2.fill_between(t_fine, pred[0], pred[2], color=\"r\", alpha=0.3)\n", " art.set_edgecolor(\"none\")\n", - " \n", - " # plot the true RA model\n", - " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color = 'k', label = \"truth\")\n", - "\n", - "\n", - "\n", - " \n", - " ax2.arrow(0.5, 0., 0, 1, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0., head_length = 0., color = 'k', ls = '--')\n", - " ax2.text(0.49, 0.9, \"Gaia\", fontsize = 18, ha='right', transform=ax2.transAxes)\n", - " ax2.text(0.51, 0.9, \"Roman\", fontsize = 18, ha='left', transform=ax2.transAxes)\n", - " ax2.arrow(0.49, 0.87, -.15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", - " ax2.arrow(0.51, 0.87, .15, 0, width = 0.001, transform=ax2.transAxes,\n", - " head_width= 0.03, head_length = .03, color ='k')\n", "\n", + " # plot the true RA model\n", + " ax2.plot(x_astrometry_orbit, y_ra_orbit.T[n], color=\"k\", label=\"truth\")\n", + "\n", + " ax2.arrow(\n", + " 0.5,\n", + " 0.0,\n", + " 0,\n", + " 1,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.0,\n", + " head_length=0.0,\n", + " color=\"k\",\n", + " ls=\"--\",\n", + " )\n", + " ax2.text(\n", + " 0.49, 0.9, \"Gaia\", fontsize=18, ha=\"right\", transform=ax2.transAxes\n", + " )\n", + " ax2.text(\n", + " 0.51, 0.9, \"Roman\", fontsize=18, ha=\"left\", transform=ax2.transAxes\n", + " )\n", + " ax2.arrow(\n", + " 0.49,\n", + " 0.87,\n", + " -0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", + " ax2.arrow(\n", + " 0.51,\n", + " 0.87,\n", + " 0.15,\n", + " 0,\n", + " width=0.001,\n", + " transform=ax2.transAxes,\n", + " head_width=0.03,\n", + " head_length=0.03,\n", + " color=\"k\",\n", + " )\n", "\n", - " \n", " ax2.legend(fontsize=9, loc=1)\n", " ax2.set_xlabel(\"time [days]\", fontsize=18)\n", - " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize = 18)\n", - " \n", - " ax3.set_ylabel(r\"[O-C]\", fontsize = 18)\n", + " ax2.set_ylabel(r\"$\\Delta \\alpha \\cos \\delta$ ['']\", fontsize=18)\n", + "\n", + " ax3.set_ylabel(r\"[O-C]\", fontsize=18)\n", " ax3.set_xlabel(\"phase [days]\", fontsize=18)\n", "\n", - " \n", - " #plot residuals\n", + " # plot residuals\n", " ax3.axhline(0.0, color=\"k\")\n", " ax3.errorbar(\n", - " x_astrometry, y_ra - other - med_ra,\n", - " yerr=y_ra_err, color = '#00257c', fmt=\".\"\n", + " x_astrometry,\n", + " y_ra - other - med_ra,\n", + " yerr=y_ra_err,\n", + " color=\"#00257c\",\n", + " fmt=\".\",\n", " )\n", - " \n", + "\n", " figAs.tight_layout()\n", - " figAs.subplots_adjust(wspace=.2)\n", - " \n", - " \n", + " figAs.subplots_adjust(wspace=0.2)\n", + "\n", " figAs.show()\n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - " \n", - "\n", - " \n", + "\n", " return None" ] }, @@ -1796,16 +1850,20 @@ "cell_type": "code", "execution_count": null, "id": "b3941ebb", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "outputs": [], "source": [ - "make_rv_astrometry_plots(simulated_data, trace, 2)\n" + "make_rv_astrometry_plots(simulated_data, trace, 2)" ] }, { "cell_type": "markdown", "id": "1f6c90cd", - "metadata": {}, + "metadata": { + "lines_to_next_cell": 2 + }, "source": [ "## And lastly, check how well our MCMC posterior median +/- 1 sigma results agree with the simulated values" ] @@ -1818,9 +1876,7 @@ "outputs": [], "source": [ "def load_posterior_params(trace, nplanets):\n", - "\n", - " parameters = ['P', 'ecc', 'tperi', 'omega', 'Omega', 'incl', 'm_planet']\n", - "\n", + " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\n", "\n", " params_earth = defaultdict(list)\n", " params_jup = defaultdict(list)\n", @@ -1835,138 +1891,139 @@ " posterior_meds = defaultdict(list)\n", " posterior_errs = defaultdict(list)\n", " for param in parameters:\n", - "\n", - "\n", " planet_med = np.median(trace.posterior[param].values[:, :, ii])\n", "\n", - " planet_quantile = [np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", - " np.quantile(trace.posterior[param].values[:, :, ii], 0.84)]\n", - "\n", - "\n", + " planet_quantile = [\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.16),\n", + " np.quantile(trace.posterior[param].values[:, :, ii], 0.84),\n", + " ]\n", "\n", " planet_err = [\n", " planet_med - planet_quantile[0],\n", - " planet_quantile[1] - planet_med \n", + " planet_quantile[1] - planet_med,\n", " ]\n", "\n", - "\n", - "\n", - " #print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", + " # print(param + \"_\" + str(ii+1) + \": \" + str(planet_med) + \" +/- \" + str(planet_err))\n", "\n", " posterior_meds[param] = np.round(planet_med, 3)\n", " posterior_errs[param] = np.round(planet_err, 3)\n", - " table_params[param].append(str(np.round(planet_med, 3)) + \" (+\" + \n", - " str(np.round(planet_err[1], 3)) + \" -\" + \n", - " str(np.round(planet_err[0], 3)) + \")\")\n", - "\n", - "\n", - "\n", - "\n", - "\n", + " table_params[param].append(\n", + " str(np.round(planet_med, 3))\n", + " + \" (+\"\n", + " + str(np.round(planet_err[1], 3))\n", + " + \" -\"\n", + " + str(np.round(planet_err[0], 3))\n", + " + \")\"\n", + " )\n", "\n", " table_params_all.append(table_params)\n", " posterior_meds_all.append(posterior_meds)\n", " posterior_errs_all.append(posterior_errs)\n", - " #print(\"\")\n", - " #print(\"\")\n", - "\n", - "\n", + " # print(\"\")\n", + " # print(\"\")\n", "\n", " return table_params_all, posterior_meds_all, posterior_errs_all\n", "\n", "\n", - "def make_comparison_figures(input_params, posterior_orbit_params, keys_orbit_params, \n", - " posterior_meds, posterior_errs, nplanets):\n", - "\n", + "def make_comparison_figures(\n", + " input_params,\n", + " posterior_orbit_params,\n", + " keys_orbit_params,\n", + " posterior_meds,\n", + " posterior_errs,\n", + " nplanets,\n", + "):\n", " import matplotlib\n", " import matplotlib.pyplot as plt\n", - " \n", - " \n", + "\n", " input_params = np.array(input_params)\n", - " input_params.T[-1] = input_params.T[-1]*m_sun\n", + " input_params.T[-1] = input_params.T[-1] * m_sun\n", "\n", " comparisons = []\n", " for ii in range(0, nplanets):\n", - " fig, axs = plt.subplots(1, 2, figsize = [13,6], gridspec_kw={'width_ratios': [1, 1]})\n", + " fig, axs = plt.subplots(\n", + " 1, 2, figsize=[13, 6], gridspec_kw={\"width_ratios\": [1, 1]}\n", + " )\n", " plt.subplots_adjust(hspace=1)\n", - " \n", + "\n", " #### add comparison plots\n", - " \n", - " \n", + "\n", " input_mass = input_params[ii][-1]\n", " input_period = input_params[ii][0]\n", - " \n", - " xs = posterior_meds[ii]['m_planet']\n", - " xs_err = np.array([posterior_errs[ii]['m_planet']]).T\n", - "\n", - " ys = posterior_meds[ii]['P']\n", - " ys_err = np.array([posterior_errs[ii]['P']]).T\n", - "\n", "\n", + " xs = posterior_meds[ii][\"m_planet\"]\n", + " xs_err = np.array([posterior_errs[ii][\"m_planet\"]]).T\n", + "\n", + " ys = posterior_meds[ii][\"P\"]\n", + " ys_err = np.array([posterior_errs[ii][\"P\"]]).T\n", + "\n", + " color = \"#1c245f\"\n", + "\n", + " axs[0].errorbar(\n", + " xs,\n", + " ys,\n", + " xerr=xs_err,\n", + " yerr=ys_err,\n", + " marker=\"o\",\n", + " color=color,\n", + " markersize=13,\n", + " label=r\"med and 1$\\sigma$ errs\",\n", + " linewidth=3,\n", + " alpha=0.7,\n", + " )\n", "\n", - " color = '#1c245f'\n", - "\n", - " axs[0].errorbar(xs, ys, xerr=xs_err, yerr = ys_err, marker = 'o', \n", - " color = color, markersize=13, label = r'med and 1$\\sigma$ errs',\n", - " linewidth = 3, alpha = 0.7)\n", - "\n", - " axs[0].axvline(x=input_mass, ymin=0, ymax=1, color = 'k', ls = '--', label = 'simulated')\n", - " axs[0].axhline(y=input_period, xmin=0, xmax=1, color = 'k', ls = '--')\n", - "\n", + " axs[0].axvline(\n", + " x=input_mass, ymin=0, ymax=1, color=\"k\", ls=\"--\", label=\"simulated\"\n", + " )\n", + " axs[0].axhline(y=input_period, xmin=0, xmax=1, color=\"k\", ls=\"--\")\n", "\n", - " axs[0].set_xlabel(r'mass [M$_{earth}$]', fontsize = 27)\n", - " axs[0].set_ylabel(r'period [days]', fontsize = 27)\n", + " axs[0].set_xlabel(r\"mass [M$_{earth}$]\", fontsize=27)\n", + " axs[0].set_ylabel(r\"period [days]\", fontsize=27)\n", "\n", + " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", "\n", - " fig.suptitle('input and posterior parameters', fontsize = 27)\n", + " axs[0].legend(fontsize=13)\n", "\n", - " axs[0].legend(fontsize = 13)\n", - " \n", - " \n", - " \n", " #### add comparison tables\n", "\n", - "\n", " sim_params = input_params[ii]\n", " post_params = posterior_orbit_params[ii]\n", - " data = np.column_stack((np.array(keys_orbit_params), \n", - " np.round(np.array(sim_params),3),\n", - " np.array(list(post_params.values()))))\n", + " data = np.column_stack(\n", + " (\n", + " np.array(keys_orbit_params),\n", + " np.round(np.array(sim_params), 3),\n", + " np.array(list(post_params.values())),\n", + " )\n", + " )\n", "\n", - " labels = ['param', 'truth', r'MCMC med +/- 1$\\sigma$']\n", + " labels = [\"param\", \"truth\", r\"MCMC med +/- 1$\\sigma$\"]\n", "\n", " df = pd.DataFrame(data, columns=labels)\n", "\n", - " table = axs[1].table(cellText=df.values, colLabels=df.columns, \n", - " colColours =[\"#D7E5F0\"]*3, loc='center', \n", - " bbox=[0,0,1,1], colWidths = [0.2, 0.2, 0.6])\n", - " #for (row, col), cell in table.get_celld().items():\n", + " table = axs[1].table(\n", + " cellText=df.values,\n", + " colLabels=df.columns,\n", + " colColours=[\"#D7E5F0\"] * 3,\n", + " loc=\"center\",\n", + " bbox=[0, 0, 1, 1],\n", + " colWidths=[0.2, 0.2, 0.6],\n", + " )\n", + " # for (row, col), cell in table.get_celld().items():\n", " # if (row == 0):\n", " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\n", "\n", - "\n", " table.auto_set_column_width(col=labels)\n", "\n", - " #ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", - " #t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", - " #t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", + " # ax.set_title('test_sim_file.csv'[:-4], fontsize=18, loc='left', fontweight =\"bold\")\n", + " # t[(np.argmin(jwst_table.values), 0)].set_facecolor(\"#53b568\")\n", + " # t[(np.argmin(jwst_table.values), 1)].set_facecolor(\"#53b568\")\n", " axs[1].set_axis_off()\n", "\n", - "\n", " table.auto_set_font_size(False)\n", " table.set_fontsize(13)\n", "\n", - " \n", - " \n", - " \n", - "\n", - "\n", - "\n", " comparisons.append(fig)\n", - " \n", - " \n", - " \n", - " \n", + "\n", " return comparisons" ] }, @@ -1980,12 +2037,29 @@ "from collections import defaultdict\n", "\n", "\n", - "table_params_all, posterior_meds_all, posterior_errs_all = \\\n", - "load_posterior_params(trace, 2)\n", - "\n", - "planet_param_keys = ['period', 'ecc', r'T$_\\mathrm{per}$', r'$\\omega$', r'$\\Omega$', 'inc', 'mass']\n", - "make_comparison_figures([planet_params[1], planet_params[0]], table_params_all, planet_param_keys, \n", - " posterior_meds_all, posterior_errs_all, 2)" + "(\n", + " table_params_all,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + ") = load_posterior_params(trace, 2)\n", + "\n", + "planet_param_keys = [\n", + " \"period\",\n", + " \"ecc\",\n", + " r\"T$_\\mathrm{per}$\",\n", + " r\"$\\omega$\",\n", + " r\"$\\Omega$\",\n", + " \"inc\",\n", + " \"mass\",\n", + "]\n", + "make_comparison_figures(\n", + " [planet_params[1], planet_params[0]],\n", + " table_params_all,\n", + " planet_param_keys,\n", + " posterior_meds_all,\n", + " posterior_errs_all,\n", + " 2,\n", + ")" ] }, {