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 new file mode 100644 index 00000000..6a29e62f --- /dev/null +++ b/docs/tutorials/astrometry_rv_case_study.ipynb @@ -0,0 +1,2103 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "fae3c154", + "metadata": {}, + "source": [ + "# Joint Astrometric and Radial Velocity Fitting" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4abd3961", + "metadata": { + "lines_to_next_cell": 2 + }, + "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__}'\")" + ] + }, + { + "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": { + "lines_to_next_cell": 2 + }, + "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", + " 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", + " 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", + "\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", + " 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", + " # 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", + " # 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 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", + "\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", + " 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", + " # 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", + "\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", + " 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", + " # 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", + " # 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", + "\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", + " rho_observed, theta_observed = rho_observed.eval(), theta_observed.eval()\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", + " 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", + " # 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", + " 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", + " # ----------- 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(\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", + " 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", + " else:\n", + " ra_sim = ra_observed + np.random.normal(0, sigma_ra, len(ra_observed))\n", + " ra_sim_sum = ra_sim\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", + " else:\n", + " dec_sim = dec_observed + np.random.normal(\n", + " 0, sigma_dec, len(dec_observed)\n", + " )\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", + " 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", + " 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 = [\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", + "]" + ] + }, + { + "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": { + "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", + "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", + "# 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", + " 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", + "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", + "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", + "[\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", + "# 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", + "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(\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", + "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)" + ] + }, + { + "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", + "\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", + " fig, ax = plt.subplots(1, 1, figsize=[18, 10], sharey=\"row\")\n", + " fig.suptitle(\"RV Signal\", fontsize=45)\n", + " ax1 = ax\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", + "\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", + " 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", + " 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(\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", + " 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": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "plot_rv_signal(\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": { + "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", + " dec_sim_sum,\n", + " times_astrometry,\n", + " times_observed_astrometry,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "99110cc0", + "metadata": { + "lines_to_next_cell": 2 + }, + "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", + " # 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", + " P = pm.Deterministic(\"P\", tt.exp(logP))\n", + "\n", + " ## wide uniform prior on t_periastron\n", + " tperi = pm.Uniform(\"tperi\", lower=0, upper=periods, shape=2)\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", + " K = pm.Deterministic(\"K\", tt.exp(logK))\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", + " # 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(\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", + " # 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", + " # 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", + " 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", + " return map_soln" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4d6faa40", + "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", + "\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", + "# 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_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", + "\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(\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", + "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": { + "lines_to_next_cell": 2 + }, + "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(\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", + " # 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_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", + " 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(\n", + " x_astrometry.min() - 500, x_astrometry.max() + 500, num=1000\n", + " )\n", + "\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", + " print(\"trying inclination = \" + str(np.degrees(inc)))\n", + " print(\"mass test val = \" + str(mass_test_vals))\n", + " print(\"------------\")\n", + "\n", + " def get_model():\n", + " with pm.Model() as model:\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", + " # 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\",\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", + " # 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", + " omega = pm.Deterministic(\"omega\", tt.arctan2(ecs[1], ecs[0]))\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", + " # 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", + " 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", + " # 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", + " # 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", + " 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", + "\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", + " # 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", + " # 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", + " # 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", + "\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(\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", + " return dec_model, ra_model\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", + "\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", + " # ADD RV MODEL\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(\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", + "\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", + " return model, map_soln\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", + " 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", + " return the_model, the_map_soln, the_logp\n", + "\n", + "\n", + "def determine_phase(P, t_periastron):\n", + " phase = (2 * np.pi * t_periastron) / P\n", + " return phase\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", + " 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,\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", + "# 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", + "\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,\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=[\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", + " plt.show()\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", + " )\n", + "\n", + " plt.show()\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", + " 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", + " 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", + " 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", + "\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" + ] + }, + { + "cell_type": "markdown", + "id": "4ca33874", + "metadata": {}, + "source": [ + "## Check the convergence diagnostics:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cfd9e0cd", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "import arviz as az\n", + "\n", + "az.summary(trace, var_names=[\"P\", \"ecc\", \"m_planet\", \"Omega\", \"omega\", \"incl\"])" + ] + }, + { + "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": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "import corner\n", + "\n", + "make_plots(trace, planet_params, 2)" + ] + }, + { + "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", + "\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", + " \"times_rv_orbit\": times_rv,\n", + " \"rv_orbit\": rv_orbit,\n", + " \"rv_orbit_sum\": rv_orbit_sum,\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", + " \"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)" + ] + }, + { + "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", + "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", + "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", + "\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", + " # 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_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", + " # 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", + " figsAs = []\n", + " figsRV = []\n", + " for n in range(0, nplanets):\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", + " 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", + " # 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", + "\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", + " # 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.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", + " 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", + " # 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", + " )\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", + " 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", + " 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", + " # 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", + "\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", + " 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.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", + "\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", + " ax1.set_xlabel(\"phase [days]\", fontsize=18)\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", + " )\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", + " # 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", + "\n", + " pred = np.percentile(\n", + " trace.posterior[\"ra_fine\"].values[:, :, :, n],\n", + " [16, 50, 84],\n", + " axis=(0, 1),\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", + " # 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.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", + " 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", + " 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", + " # 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", + " )\n", + "\n", + " figAs.tight_layout()\n", + " figAs.subplots_adjust(wspace=0.2)\n", + "\n", + " figAs.show()\n", + "\n", + " return None" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b3941ebb", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "make_rv_astrometry_plots(simulated_data, trace, 2)" + ] + }, + { + "cell_type": "markdown", + "id": "1f6c90cd", + "metadata": { + "lines_to_next_cell": 2 + }, + "source": [ + "## And lastly, check how well our MCMC posterior 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", + " parameters = [\"P\", \"ecc\", \"tperi\", \"omega\", \"Omega\", \"incl\", \"m_planet\"]\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", + " 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", + "\n", + " planet_err = [\n", + " planet_med - planet_quantile[0],\n", + " planet_quantile[1] - planet_med,\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", + "\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", + " 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", + " import matplotlib\n", + " import matplotlib.pyplot as plt\n", + "\n", + " input_params = np.array(input_params)\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", + " plt.subplots_adjust(hspace=1)\n", + "\n", + " #### add comparison plots\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", + " 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", + " 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", + "\n", + " fig.suptitle(\"input and posterior parameters\", fontsize=27)\n", + "\n", + " axs[0].legend(fontsize=13)\n", + "\n", + " #### add comparison tables\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", + "\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", + " # if (row == 0):\n", + " # cell.set_text_props(fontproperties=FontProperties(weight='bold'))\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", + " table.auto_set_font_size(False)\n", + " table.set_fontsize(13)\n", + "\n", + " comparisons.append(fig)\n", + "\n", + " return comparisons" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "552e3f12", + "metadata": {}, + "outputs": [], + "source": [ + "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", + ")" + ] + }, + { + "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 +} 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",