Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Update case studies to pymc v5 #78

Open
wants to merge 15 commits into
base: main
Choose a base branch
from
Open

Update case studies to pymc v5 #78

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
cf2eca1
Update RV tutorial to PyMC v5
vandalt Mar 2, 2024
a9d05c2
Update transit tutorial to PyMC v5
vandalt Mar 2, 2024
1181a14
Update astrometric model to PyMC v5
vandalt Mar 2, 2024
46449fc
Update stellar variability tutorial to PyMC v5
vandalt Mar 2, 2024
13fbfca
Remove leftover todo in astrometric notebok (and format)
vandalt Mar 2, 2024
1f53767
Update joint RV+Transit notebook to PyMC v5
vandalt Mar 2, 2024
7274646
Remove leftover todo in joint RV+Transit notebok
vandalt Mar 2, 2024
d68adc5
Update TESS tutorial to PyMC v5
vandalt Mar 2, 2024
9da14fc
Update quick TESS notebook to PyMC v5
vandalt Mar 2, 2024
089f41a
Update TTV tutorial to PyMC v5
vandalt Mar 2, 2024
4dd60f7
Update multi-instrument RV tutorial to PyMC v5
vandalt Mar 3, 2024
75325fe
Update transit+activity tutorial to PyMC v5
vandalt Mar 5, 2024
e53fbd0
Update eclipsing binary tutorial to PyMC v5
vandalt Mar 5, 2024
e076871
Update multi-instrument light curve tutorial to PyMC v5
vandalt Mar 6, 2024
6040ac6
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Mar 6, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
73 changes: 42 additions & 31 deletions docs/tutorials/astrometric.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,7 +93,7 @@
"astro_data[\"rho_err\"][astro_data[\"rho_err\"].mask == True] = 0.01\n",
"astro_data[\"PA_err\"][astro_data[\"PA_err\"].mask == True] = 1.0\n",
"\n",
"# Convert all masked frames to be raw np arrays, since theano has issues with astropy masked columns\n",
"# Convert all masked frames to be raw np arrays, since pytensor has issues with astropy masked columns\n",
"rho_data = np.ascontiguousarray(astro_data[\"rho\"], dtype=float) # arcsec\n",
"rho_err = np.ascontiguousarray(astro_data[\"rho_err\"], dtype=float)\n",
"\n",
Expand Down Expand Up @@ -192,11 +192,11 @@
"metadata": {},
"outputs": [],
"source": [
"import pymc3 as pm\n",
"import pymc3_ext as pmx\n",
"import pymc as pm\n",
"import pymc_ext as pmx\n",
"\n",
"from aesara_theano_fallback import aesara as theano\n",
"import aesara_theano_fallback.tensor as tt\n",
"import pytensor\n",
"import pytensor.tensor as pt\n",
"\n",
"import exoplanet as xo\n",
"\n",
Expand Down Expand Up @@ -230,7 +230,7 @@
"# The position functions take an optional argument parallax to convert from\n",
"# physical units back to arcseconds\n",
"t = np.linspace(T0 - P, T0 + P, num=200) # days\n",
"rho, theta = theano.function([], orbit.get_relative_angles(t, parallax))()\n",
"rho, theta = pytensor.function([], orbit.get_relative_angles(t, parallax))()\n",
"\n",
"# Plot the orbit\n",
"fig, ax = plt.subplots(nrows=1, figsize=(4, 4))\n",
Expand Down Expand Up @@ -268,7 +268,7 @@
"id": "ecd7ec33",
"metadata": {},
"source": [
"Now that we have an initial orbit, we can set the model up using PyMC3 to do inference."
"Now that we have an initial orbit, we can set the model up using PyMC to do inference."
]
},
{
Expand All @@ -293,43 +293,47 @@
" else:\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=100)(\n",
" \"m_plx\", mu=parallax[0], sd=parallax[1], testval=parallax[0]\n",
" m_plx = pm.Truncated(\n",
" \"m_plx\",\n",
" pm.Normal.dist(mu=parallax[0], sigma=parallax[1]),\n",
" lower=0,\n",
" upper=100,\n",
" initval=parallax[0],\n",
" )\n",
" plx = pm.Deterministic(\"plx\", 1e-3 * m_plx)\n",
"\n",
" a_ang = pm.Uniform(\"a_ang\", 0.1, 1.0, testval=0.324)\n",
" a_ang = pm.Uniform(\"a_ang\", 0.1, 1.0, initval=0.324)\n",
" a = pm.Deterministic(\"a\", a_ang / plx)\n",
"\n",
" # We expect the period to be somewhere in the range of 25 years,\n",
" # so we'll set a broad prior on logP\n",
" logP = pm.Normal(\n",
" \"logP\", mu=np.log(25 * yr), sd=10.0, testval=np.log(28.8 * yr)\n",
" \"logP\", mu=np.log(25 * yr), sigma=10.0, initval=np.log(28.8 * yr)\n",
" )\n",
" P = pm.Deterministic(\"P\", tt.exp(logP))\n",
" P = pm.Deterministic(\"P\", pt.exp(logP))\n",
"\n",
" # For astrometric-only fits, it's generally better to fit in\n",
" # p = (Omega + omega)/2 and m = (Omega - omega)/2 instead of omega and Omega\n",
" # directly\n",
" omega0 = 251.6 * deg - np.pi\n",
" Omega0 = 159.6 * deg\n",
" p = pmx.Angle(\"p\", testval=0.5 * (Omega0 + omega0))\n",
" m = pmx.Angle(\"m\", testval=0.5 * (Omega0 - omega0))\n",
" p = pmx.angle(\"p\", initval=0.5 * (Omega0 + omega0))\n",
" m = pmx.angle(\"m\", initval=0.5 * (Omega0 - omega0))\n",
" omega = pm.Deterministic(\"omega\", p - m)\n",
" Omega = pm.Deterministic(\"Omega\", p + m)\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 periastron\n",
" # passage directly\n",
" phase = pmx.Angle(\"phase\", testval=0.0)\n",
" phase = pmx.angle(\"phase\", initval=0.0)\n",
" tperi = pm.Deterministic(\"tperi\", T0 + P * phase / (2 * np.pi))\n",
"\n",
" # Geometric uniform prior on cos(incl)\n",
" cos_incl = pm.Uniform(\n",
" \"cos_incl\", lower=-1, upper=1, testval=np.cos(96.0 * deg)\n",
" \"cos_incl\", lower=-1, upper=1, initval=np.cos(96.0 * deg)\n",
" )\n",
" incl = pm.Deterministic(\"incl\", tt.arccos(cos_incl))\n",
" ecc = pm.Uniform(\"ecc\", lower=0.0, upper=1.0, testval=0.798)\n",
" incl = pm.Deterministic(\"incl\", pt.arccos(cos_incl))\n",
" ecc = pm.Uniform(\"ecc\", lower=0.0, upper=1.0, initval=0.798)\n",
"\n",
" # Set up the orbit\n",
" orbit = xo.orbits.KeplerianOrbit(\n",
Expand All @@ -351,32 +355,39 @@
"\n",
" # Add jitter terms to both separation and position angle\n",
" log_rho_s = pm.Normal(\n",
" \"log_rho_s\", mu=np.log(np.median(rho_err)), sd=2.0\n",
" \"log_rho_s\", mu=np.log(np.median(rho_err)), sigma=2.0\n",
" )\n",
" log_theta_s = pm.Normal(\n",
" \"log_theta_s\", mu=np.log(np.median(theta_err)), sd=2.0\n",
" \"log_theta_s\", mu=np.log(np.median(theta_err)), sigma=2.0\n",
" )\n",
" rho_tot_err = tt.sqrt(rho_err**2 + tt.exp(2 * log_rho_s))\n",
" theta_tot_err = tt.sqrt(theta_err**2 + tt.exp(2 * log_theta_s))\n",
" rho_tot_err = pt.sqrt(rho_err**2 + pt.exp(2 * log_rho_s))\n",
" theta_tot_err = pt.sqrt(theta_err**2 + pt.exp(2 * log_theta_s))\n",
"\n",
" # define the likelihood function, e.g., a Gaussian on both rho and theta\n",
" pm.Normal(\"rho_obs\", mu=rho_model, sd=rho_tot_err, observed=rho_data)\n",
" pm.Normal(\n",
" \"rho_obs\", mu=rho_model, sigma=rho_tot_err, observed=rho_data\n",
" )\n",
"\n",
" # We want to be cognizant of the fact that theta wraps so the following is equivalent to\n",
" # pm.Normal(\"obs_theta\", mu=theta_model, observed=theta_data, sd=theta_tot_err)\n",
" # but takes into account the wrapping. Thanks to Rob de Rosa for the tip.\n",
" theta_diff = tt.arctan2(\n",
" tt.sin(theta_model - theta_data), tt.cos(theta_model - theta_data)\n",
" theta_diff = pt.arctan2(\n",
" pt.sin(theta_model - theta_data), pt.cos(theta_model - theta_data)\n",
" )\n",
" pm.Normal(\n",
" \"theta_obs\",\n",
" mu=theta_diff,\n",
" sigma=theta_tot_err,\n",
" observed=np.zeros_like(theta_data),\n",
" )\n",
" pm.Normal(\"theta_obs\", mu=theta_diff, sd=theta_tot_err, observed=0.0)\n",
"\n",
" # Set up predicted orbits for later plotting\n",
" rho_dense, theta_dense = orbit.get_relative_angles(t_fine, plx)\n",
" rho_save = pm.Deterministic(\"rho_save\", rho_dense)\n",
" theta_save = pm.Deterministic(\"theta_save\", theta_dense)\n",
"\n",
" # Optimize to find the initial parameters\n",
" map_soln = model.test_point\n",
" map_soln = model.initial_point()\n",
" map_soln = pmx.optimize(map_soln, vars=[log_rho_s, log_theta_s])\n",
" map_soln = pmx.optimize(map_soln, vars=[phase])\n",
" map_soln = pmx.optimize(map_soln, vars=[p, m, ecc])\n",
Expand Down Expand Up @@ -453,10 +464,10 @@
"source": [
"np.random.seed(1234)\n",
"with model:\n",
" trace = pmx.sample(\n",
" trace = pmx.utils.sample(\n",
" tune=1000,\n",
" draws=1000,\n",
" start=map_soln,\n",
" initvals=map_soln,\n",
" cores=2,\n",
" chains=2,\n",
" target_accept=0.9,\n",
Expand Down Expand Up @@ -596,10 +607,10 @@
"source": [
"np.random.seed(5432)\n",
"with plx_model:\n",
" plx_trace = pmx.sample(\n",
" plx_trace = pmx.utils.sample(\n",
" tune=1000,\n",
" draws=1000,\n",
" start=plx_map_soln,\n",
" initvals=plx_map_soln,\n",
" cores=2,\n",
" chains=2,\n",
" target_accept=0.9,\n",
Expand Down
Loading
Loading