<jupyter_start><jupyter_text>This notebook demonstrates how to use kepler_kinematics to calculate stellar velocities<jupyter_code>%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
from astropy.io import fits
import pymc3 as pm

import corner

import kepler_kinematics as kek<jupyter_output><empty_output><jupyter_text>Let's load the Gaia-Kepler crossmatched catalog.<jupyter_code>with fits.open("../../data/kepler_dr2_1arcsec.fits") as data:
    gaia = pd.DataFrame(data[1].data, dtype="float64")
gaia_rv = gaia.iloc[gaia.radial_velocity.values != 0.]
# df = pd.read_csv("../kepler_kinematics/data.csv")
# gaia_rv = df.iloc[np.isfinite(df.rv.values)]<jupyter_output><empty_output><jupyter_text>Select a star.<jupyter_code>np.random.seed(1)
star_index = np.random.randint(0, len(gaia_rv))
print(star_index)
df = gaia_rv.iloc[star_index]
print(int(df.kepid))<jupyter_output>29733
12218729<jupyter_text>First, let's calculate the 3D velocities of the star using its radial velocity.<jupyter_code># for i, row in enumerate(df):
xyz, vxyz = kek.simple_calc_vxyz(df["ra"], df["dec"],
#                                  df["r_med_photogeo"]*1e-3,
                                 df["r_est"]*1e-3,
                                 df["pmra"], df["pmdec"],
                                 df["radial_velocity"])
#                                  df["rv"])


vx, vy, vz = vxyz
x, y, z = xyz
# d = df["r_med_photogeo"] * 1e-3
d = df["r_est"] * 1e-3<jupyter_output><empty_output><jupyter_text>Now let's infer the 3D velocities of the star. First we'll create arrays of position, parallax, and proper motion.<jupyter_code>pos = [float(df["ra"]), float(df["dec"]), float(df["parallax"])]
pos_err = [float(df["ra_error"]), float(df["dec_error"]),
           float(df["parallax_error"])]
proper = [float(df["pmra"]), float(df["pmdec"])]
proper_err = [float(df["pmra_error"]), float(df["pmdec_error"])]<jupyter_output><empty_output><jupyter_text>Get the prior that's built-in for Kepler stars.<jupyter_code>mu, cov = kek.get_prior()<jupyter_output><empty_output><jupyter_text>Run the velocity inference code.<jupyter_code>trace = kek.run_pymc3_model(pos, pos_err, proper, proper_err, mu, cov)<jupyter_output>optimizing logp for variables: [vxyzD]<jupyter_text>Pull out the results from the samples.<jupyter_code>import arviz as az

az.summary(trace)
# trace.posterior.to_dataframe()
print(np.median(trace.posterior["D"].values))
print(np.median(trace.posterior["vxyz"].values, axis=(0, 1)))
print(np.shape(trace.posterior["vxyzD"].values))
vsamples = np.reshape(trace.posterior["vxyzD"].values, (4000, 4))
print(np.shape(vsamples))
print(np.shape(np.cov(vsamples.T)))
print((np.cov(vsamples.T)))
cov_mat = np.cov(vsamples.T)
vxvx_covar, vxvy_covar, vxvz_covar, vxlnd_covar = cov_mat[:, 0]
vyvx_covar, vyvy_covar, vyvz_covar, vylnd_covar = cov_mat[:, 1]
vzvx_covar, vzvy_covar, vzvz_covar, vzlnd_covar = cov_mat[:, 2]
lndvx_covar, lndvy_covar, lndvz_covar, lndlnd_covar = cov_mat[:, 3]

print(vxvx_covar, vxvy_covar, vxvz_covar, vxlnd_covar,
      vyvy_covar, vyvz_covar, vylnd_covar,
      vzvz_covar, vzlnd_covar,
      lndlnd_covar)
a = np.linspace(0, 10, 10)
print(np.shape(np.vstack((a, a))))
plotpar = {'axes.labelsize': 25,
           'font.size': 12,
           'legend.fontsize': 15,
           'xtick.labelsize': 15,
           'ytick.labelsize': 15,
           'text.usetex': True}
plt.rcParams.update(plotpar)

labels = ["$\mathrm{V_x~[kms^{-1}]}$", "$\mathrm{V_y~[kms^{-1}]}$", "$\mathrm{V_z~[kms^{-1}]}$",
          "$\mathrm{\ln(D/kpc)}$"]
fig = corner.corner(vsamples, labels=labels, truths=[vx.value, vy.value, vz.value, np.log(d)], dpi=200);
# plt.title(f"KIC {format(int(df.kepid))}")
plt.savefig("../plots/posterior")
plt.savefig("../paper/posterior.pdf")
print(int(df.kepid))
d
flat_samples = pm.trace_to_dataframe(trace)

params_inferred = np.median(flat_samples, axis=0)
upper = np.percentile(flat_samples, 84, axis=0)
lower = np.percentile(flat_samples, 16, axis=0)
errp = upper - params_inferred
errm = params_inferred - lower
std = np.std(flat_samples, axis=0)

results = pd.DataFrame(dict({"kepid": df["kepid"],
                             "vx": params_inferred[0], "vx_errp": errp[0], "vx_errm": errm[0], "vx_err": std[0],
                             "vy": params_inferred[1], "vy_errp": errp[1], "vy_errm": errm[1], "vy_err": std[1],
                             "vz": params_inferred[2], "vz_errp": errp[2], "vz_errm": errm[2], "vz_err": std[2],
                             "lndistance": params_inferred[3], "lndistance_errp": errp[3],
                             "lndistance_errm": errm[3], "lndistance_err": std[3]
                            }), index=[0])
print(f"{vx:.1f}, {float(results.vx.values):.1f} +/- {float(results.vx_err.values):.1f} km / s")
print(f"{vy:.1f}, {float(results.vy.values):.1f} +/- {float(results.vy_err.values):.1f} km / s")
print(f"{vz:.1f}, {float(results.vz.values):.1f} +/- {float(results.vz_err.values):.1f} km / s")
print(f"{vx:.1f}, {float(results.vx.values):.1f} +/- {float(results.vx_err.values):.1f} km / s")
print(f"{vy:.1f}, {float(results.vy.values):.1f} +/- {float(results.vy_err.values):.1f} km / s")
print(f"{vz:.1f}, {float(results.vz.values):.1f} +/- {float(results.vz_err.values):.1f} km / s")<jupyter_output>50.4 km / s, 52.8 +/- 8.7 km / s
215.2 km / s, 220.7 +/- 22.7 km / s
-0.3 km / s, 0.7 +/- 4.2 km / s<jupyter_text>Test with the new data file.<jupyter_code>df = pd.read_csv("../kepler_kinematics/data.csv")
gaia_rv = df.iloc[np.isfinite(df.rv.values)]
df = gaia_rv.iloc[0]
print(int(df.kepid))
xyz, vxyz = kek.simple_calc_vxyz(df["ra"], df["dec"],
                                 df["r_med_photogeo"]*1e-3,
                                 df["pmra"], df["pmdec"],
                                 df["rv"])

vx, vy, vz = vxyz
x, y, z = xyz
d = df["r_med_photogeo"] * 1e-3
pos = [float(df["ra"]), float(df["dec"]), float(df["parallax"])]
pos_err = [float(df["ra_error"]), float(df["dec_error"]),
           float(df["parallax_error"])]
proper = [float(df["pmra"]), float(df["pmdec"])]
proper_err = [float(df["pmra_error"]), float(df["pmdec_error"])]
mu, cov = kek.get_prior()
trace = kek.run_pymc3_model(pos, pos_err, proper, proper_err, mu, cov)
plotpar = {'axes.labelsize': 25,
           'font.size': 12,
           'legend.fontsize': 15,
           'xtick.labelsize': 15,
           'ytick.labelsize': 15,
           'text.usetex': True}
plt.rcParams.update(plotpar)

labels = ["$\mathrm{V_x~[kms^{-1}]}$", "$\mathrm{V_y~[kms^{-1}]}$", "$\mathrm{V_z~[kms^{-1}]}$",
          "$\mathrm{\ln(D/kpc)}$"]
fig = corner.corner(vsamples, labels=labels, truths=[df.vx, df.vy, df.vz,
                                                     np.log(df.r_med_photogeo*1e-3)], dpi=200);

axes = np.array(fig.axes).reshape((4, 4))
ax = axes[0, 0]
ax.set_title(f"KIC {format(int(df.kepid))}")
df.vx, df.vy, df.vz, np.log(df.r_med_photogeo*1e-3)
res = pd.read_csv("../data/edr3_all.csv")
np.array([k for k in res.keys()])
resdat = pd.merge(res, gaia_rv, on="kepid")
plt.plot(resdat.vx_inferred, resdat.vx-resdat.vx_inferred, ".", alpha=.1, ms=1)
plt.plot(resdat.vx_inferred[0], resdat.vx[0]-resdat.vx_inferred[0], ".")
plt.ylim(-10, 10)
plt.axhline(0, color="k")<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.19144267690619857, 'fraction_numerical': 0.023861766319253977, 'mean_word_length': 3.7697841726618706, 'pattern_counts': {'":': 17, '<': 34, '<?xml version=': 0, '>': 34, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '9', 'dir': 'jupyter-structured-clean-dedup', 'id': '446715', 'license': 'MIT', 'n_tokens_mistral': 3146, 'n_tokens_neox': 2949, 'n_words': 481, 'path': 'notebooks/Demo_Notebook.ipynb', 'repo_name': 'RuthAngus/kepler_kinematics', 'size': '7292'}