scripts/precision_four_panel.py
"""
precision_four_panel.py
-----------------------
Plot a Figueira et al. (2016) Figure 1 like plot.
"""
import argparse
import sys
from os.path import join
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from eniric import config
from eniric.utilities import rv_cumulative, rv_cumulative_full
matplotlib.use("Agg")
def load_dataframe(filename):
"""Load in precision file, clean up spaces in csv.
Parameters
----------
filename: str
Name of phoenix_precision.py output.
Returns
-------
df: pandas.DataFrame
DataFrame of data.
"""
df = pd.read_csv(precision_file)
# Temp, logg, [Fe/H], Alpha, Band, Resolution, vsini, Sampling, Quality, Cond. 1, Cond. 2, Cond. 3, correct flag
df.columns = df.columns.str.strip()
df.Band = df.Band.str.strip()
df.Resolution = df.Resolution.str.strip()
df.Quality = df.Quality.astype(float)
df.Temp = df.Temp.astype(float)
df.logg = df.logg.astype(float)
df["[Fe/H]"] = df["[Fe/H]"].astype(float)
df.Alhpa = df.Alpha.astype(float)
df.vsini = df.vsini.astype(float)
df.Sampling = df.Sampling.astype(float)
df["Cond. 1"] = df["Cond. 1"].astype(float)
df["Cond. 2"] = df["Cond. 2"].astype(float)
df["Cond. 3"] = df["Cond. 3"].astype(float)
df["correct flag"] = df["correct flag"].astype(bool)
return df
def plot_precision(
precision_file, teffs=None, logg=4.5, fe_h=0.0, vsini=1.0, sampling=3
):
"""Plot precision 4 panel with RV precision.
Saves figure to ``plots/``.
Parameters
----------
precision_file: str
Name of phoenix_precision.py output.
teffs: List of int or None
Stellar temperatures. Default is [3900, 3500, 2800, 2600].
logg: int
Stellar Logg. Default is 4.5.
fe_h: int
Stellar metallicity. Default is 0.0.
vsini: float
Rotational velocity. Default is 1.0.
sampling:
Spectral sampling. Default is 3.
"""
if teffs is None:
# Default teffs
teffs = [3900, 3500, 2800, 2600]
assert len(teffs) == 4
df = load_dataframe(precision_file)
filter_dict = {"logg": logg, "[Fe/H]": fe_h, "vsini": vsini, "Sampling": sampling}
df = filter_df(
df, filter_dict, drop_list=["Alpha", "[Fe/H]", "correct flag", "Quality"]
)
fig, axes = plt.subplots(2, 2)
ax = axes.flatten()
df_selected = df[df.Resolution.isin(["60k", "80k", "100k"])]
df_selected = df_selected[df_selected.Temp.isin(teffs)]
maximums = []
minimums = []
for ii, temp in enumerate(teffs):
# This entry
df_ii = df_selected[df_selected["Temp"] == temp]
df_ii_60k = df_ii[df_ii["Resolution"].str.strip() == "60k"]
df_ii_80k = df_ii[df_ii["Resolution"].str.strip() == "80k"]
df_ii_100k = df_ii[df_ii["Resolution"].str.strip() == "100k"]
df_ii_60k = df_ii_60k.set_index("Band")
df_ii_60k = df_ii_60k.reindex(["Z", "Y", "J", "H", "K"])
df_ii_80k = df_ii_80k.set_index("Band")
df_ii_80k = df_ii_80k.reindex(["Z", "Y", "J", "H", "K"])
df_ii_100k = df_ii_100k.set_index("Band")
df_ii_100k = df_ii_100k.reindex(["Z", "Y", "J", "H", "K"])
maximums.append(
np.max(
[
df_ii_60k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
df_ii_80k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
df_ii_100k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
]
)
)
minimums.append(
np.min(
[
df_ii_60k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
df_ii_80k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
df_ii_100k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
]
)
)
ax[ii].fill_between(
df_ii_60k.index,
df_ii_60k["Cond. 2"].values,
df_ii_60k["Cond. 3"].values,
color="b",
alpha=0.2,
)
ax[ii].fill_between(
df_ii_80k.index,
df_ii_80k["Cond. 2"].values,
df_ii_80k["Cond. 3"].values,
color="g",
alpha=0.2,
)
ax[ii].fill_between(
df_ii_100k.index,
df_ii_100k["Cond. 2"].values,
df_ii_100k["Cond. 3"].values,
color="r",
alpha=0.2,
)
ax[ii].plot(
df_ii_60k.index, df_ii_60k["Cond. 1"].values, color="b", linestyle="--"
) # lim
ax[ii].plot(
df_ii_80k.index, df_ii_80k["Cond. 1"].values, color="g", linestyle="--"
) # lim
ax[ii].plot(
df_ii_100k.index, df_ii_100k["Cond. 1"].values, color="r", linestyle="--"
) # lim
ax[ii].scatter(
df_ii_60k.index,
df_ii_60k["Cond. 2"].values,
marker="^",
color="b",
alpha=0.4,
)
ax[ii].scatter(
df_ii_60k.index,
df_ii_60k["Cond. 3"].values,
marker="o",
color="b",
alpha=0.4,
)
ax[ii].scatter(
df_ii_80k.index,
df_ii_80k["Cond. 3"].values,
marker="^",
color="g",
alpha=0.4,
)
ax[ii].scatter(
df_ii_80k.index,
df_ii_80k["Cond. 2"].values,
marker="o",
color="g",
alpha=0.4,
)
ax[ii].scatter(
df_ii_100k.index,
df_ii_100k["Cond. 3"].values,
marker="^",
color="r",
alpha=0.4,
)
ax[ii].scatter(
df_ii_100k.index,
df_ii_100k["Cond. 2"].values,
marker="o",
color="r",
alpha=0.4,
)
# Set limits ticks and labels
ymax = np.max(maximums)
ymin = np.min(minimums)
delta_y = ymax - ymin
band_size = len(df_ii_60k.index)
for jj in range(4):
ax[jj].text(0, ymax, "{} K".format(teffs[jj]), size=14)
ax[jj].set_ylim(ymin - 0.11 * delta_y, ymax + 0.11 * delta_y)
ax[jj].set_xlim(-0.5, band_size - 0.5)
ax[jj].tick_params(axis="both", which="major", labelsize=12)
# ticks and labels
if (jj == 2) or (jj == 3):
ax[jj].set_xlabel("Bands", fontsize=12)
if (jj == 1) or (jj == 3):
ax[jj].set_yticklabels([])
if (jj == 0) or (jj == 1):
ax[jj].set_xticklabels([])
fig.text(
0.04,
0.5,
r"Precision [m/s]",
ha="center",
va="center",
rotation="vertical",
size=12,
)
fig.subplots_adjust(hspace=0, wspace=0, bottom=0.12, top=0.95, right=0.95)
fig.savefig("plots/precision_logg{0}_feh_{1}.pdf".format(logg, fe_h))
fig.savefig("plots/precision_logg{0}_feh_{1}.png".format(logg, fe_h), dpi=400)
def filter_df(df, filter_dict, drop_list=None):
"""Filter DataFrame by dictionary of key and values."""
for key, val in filter_dict.items():
df = df[df[key] == val]
if drop_list is not None:
df = df.drop(columns=drop_list)
return df
def cumulative_df(df, full_cum=False):
"""Calculated cumulative RV precision across bands.
The precision of "Z", "ZY", "ZYJ", "ZYJH", "ZYJHK" bands.
Parameters
----------
df: pandas.DataFrame
DataFrame.
full_cum: bool
Include "YJHK", "JHK", "HK", "K" grouping also. Default is False.
"""
bands = df.index
assert all(bands == ["Z", "Y", "J", "H", "K"]), bands
if full_cum:
cum_bands = ["Z", "ZY", "ZYJ", "ZYJH", "ZYJHK", "YJHK", "JHK", "HK", "K"]
cum_dict = {
"Band": cum_bands,
"Cond. 1": rv_cumulative_full(df["Cond. 1"]),
"Cond. 2": rv_cumulative_full(df["Cond. 2"]),
"Cond. 3": rv_cumulative_full(df["Cond. 3"]),
}
else:
cum_bands = ["Z", "ZY", "ZYJ", "ZYJH", "ZYJHK"]
cum_dict = {
"Band": cum_bands,
"Cond. 1": rv_cumulative(df["Cond. 1"], single=True),
"Cond. 2": rv_cumulative(df["Cond. 2"], single=True),
"Cond. 3": rv_cumulative(df["Cond. 3"], single=True),
}
cum_df = pd.DataFrame(cum_dict)
cum_df = cum_df.set_index("Band")
cum_df = cum_df.reindex(cum_bands)
return cum_df
def cumulative_plot(
precision_file,
teffs=None,
logg=4.5,
fe_h=0.0,
vsini=1.0,
sampling=3,
full_cum=False,
):
"""RV precision with cumulative bands.
full_cum: bool
Cumlative over entire range [ "Z","ZY", "ZYJ", "ZYJH", "ZYJHK","YJHK", "JHK","HK","K"]
Saves figure to ``plots/``.
Parameters
----------
precision_file: str
Name of phoenix_precision.py output.
teffs: List of int or None
Stellar temperatures. Default is [3900, 3500, 2800, 2600].
logg: int
Stellar Logg. Default is 4.5.
fe_h: int
Stellar metallicity. Default is 0.0.
vsini: float
Rotational velocity. Default is 1.0.
sampling:
Spectral sampling. Default is 3.
full_cum: bool
Cumulative over entire range. Default is False.
"""
if teffs is None:
# Default values
teffs = [3900, 3500, 2800, 2600]
assert len(teffs) == 4
df = load_dataframe(precision_file)
filter_dict = {"logg": logg, "[Fe/H]": fe_h, "vsini": vsini, "Sampling": sampling}
df = filter_df(
df, filter_dict, drop_list=["Alpha", "[Fe/H]", "correct flag", "Quality"]
)
fig, axes = plt.subplots(2, 2)
ax = axes.flatten()
df_selected = df[df.Resolution.isin(["60k", "80k", "100k"])]
df_selected = df_selected[df_selected.Temp.isin(teffs)]
maximums = []
minimums = []
for ii, temp in enumerate(teffs):
# This entry
df_ii = df_selected[df_selected["Temp"] == temp]
df_ii_60k = df_ii[df_ii["Resolution"].str.strip() == "60k"]
df_ii_80k = df_ii[df_ii["Resolution"].str.strip() == "80k"]
df_ii_100k = df_ii[df_ii["Resolution"].str.strip() == "100k"]
df_ii_60k = df_ii_60k.set_index("Band")
df_ii_60k = df_ii_60k.reindex(["Z", "Y", "J", "H", "K"])
df_ii_80k = df_ii_80k.set_index("Band")
df_ii_80k = df_ii_80k.reindex(["Z", "Y", "J", "H", "K"])
df_ii_100k = df_ii_100k.set_index("Band")
df_ii_100k = df_ii_100k.reindex(["Z", "Y", "J", "H", "K"])
# Cumulative
df_ii_60k = cumulative_df(df_ii_60k, full_cum=full_cum)
df_ii_80k = cumulative_df(df_ii_80k, full_cum=full_cum)
df_ii_100k = cumulative_df(df_ii_100k, full_cum=full_cum)
maximums.append(
np.max(
[
df_ii_60k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
df_ii_80k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
df_ii_100k[["Cond. 1", "Cond. 2", "Cond. 3"]].max(),
]
)
)
minimums.append(
np.min(
[
df_ii_60k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
df_ii_80k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
df_ii_100k[["Cond. 1", "Cond. 2", "Cond. 3"]].min(),
]
)
)
ax[ii].fill_between(
df_ii_60k.index,
df_ii_60k["Cond. 2"].values,
df_ii_60k["Cond. 3"].values,
color="b",
alpha=0.2,
)
ax[ii].fill_between(
df_ii_80k.index,
df_ii_80k["Cond. 2"].values,
df_ii_80k["Cond. 3"].values,
color="g",
alpha=0.2,
)
ax[ii].fill_between(
df_ii_100k.index,
df_ii_100k["Cond. 2"].values,
df_ii_100k["Cond. 3"].values,
color="r",
alpha=0.2,
)
ax[ii].plot(
df_ii_60k.index, df_ii_60k["Cond. 1"].values, color="b", linestyle="--"
) # lim
ax[ii].plot(
df_ii_80k.index, df_ii_80k["Cond. 1"].values, color="g", linestyle="--"
) # lim
ax[ii].plot(
df_ii_100k.index, df_ii_100k["Cond. 1"].values, color="r", linestyle="--"
) # lim
ax[ii].scatter(
df_ii_60k.index,
df_ii_60k["Cond. 2"].values,
marker="^",
color="b",
alpha=0.4,
)
ax[ii].scatter(
df_ii_60k.index,
df_ii_60k["Cond. 3"].values,
marker="o",
color="b",
alpha=0.4,
)
ax[ii].scatter(
df_ii_80k.index,
df_ii_80k["Cond. 3"].values,
marker="^",
color="g",
alpha=0.4,
)
ax[ii].scatter(
df_ii_80k.index,
df_ii_80k["Cond. 2"].values,
marker="o",
color="g",
alpha=0.4,
)
ax[ii].scatter(
df_ii_100k.index,
df_ii_100k["Cond. 3"].values,
marker="^",
color="r",
alpha=0.4,
)
ax[ii].scatter(
df_ii_100k.index,
df_ii_100k["Cond. 2"].values,
marker="o",
color="r",
alpha=0.4,
)
ax[ii].fill_between(
df_ii_100k.index,
df_ii_100k["Cond. 2"].values,
df_ii_100k["Cond. 3"].values,
color="r",
alpha=0.2,
)
# Set limits ticks and labels
ymax = np.min([np.max(maximums), 10]) # Set limit to 20 km/s
ymin = np.min(minimums)
delta_y = ymax - ymin
band_size = len(df_ii_60k.index)
for jj in range(4):
ax[jj].text(0, ymax, "{} K".format(teffs[jj]), size=14)
ax[jj].set_ylim(ymin - 0.1 * delta_y, ymax + 0.15 * delta_y)
ax[jj].set_xlim(-0.5, band_size - 0.5)
ax[jj].tick_params(axis="both", which="major", labelsize=12)
# ticks and labels
if (jj == 2) or (jj == 3):
ax[jj].set_xlabel("Bands", fontsize=12)
if full_cum:
ax[jj].tick_params(axis="x", labelrotation=40)
else:
ax[jj].tick_params(axis="x", labelrotation=25)
if (jj == 1) or (jj == 3):
ax[jj].set_yticklabels([])
if (jj == 0) or (jj == 1):
ax[jj].set_xticklabels([])
fig.text(
0.04,
0.5,
r"Precision [m/s]",
ha="center",
va="center",
rotation="vertical",
size=12,
)
fig.subplots_adjust(hspace=0, wspace=0, bottom=0.17, top=0.95, right=0.95)
fname = "plots/cummulative_precision_logg{0}_feh_{1}_{2}".format(
logg, fe_h, full_cum
)
fig.savefig(fname + ".pdf")
fig.savefig(fname + ".png", dpi=400)
default_file = join(
config.pathdir, config.paths["precision_results"], "precision_results.csv"
)
def _parser():
"""Take care of all the argparse stuff."""
parser = argparse.ArgumentParser(
description="Create a four panel RV relative precision plot."
)
parser.add_argument(
"precision_file",
help="Precision result csv to use. Default is set with the config.yaml file. Currently {0}".format(
default_file
),
type=str,
default=default_file,
)
parser.add_argument(
"-t",
"--temperatures",
nargs=4,
help="Temperatures to display in Kelvin. Default is [3900, 3500, 2800, 2600].",
type=int,
default=[3900, 3500, 2800, 2600],
)
return parser.parse_args()
if __name__ == "__main__":
args = _parser()
precision_file = args.precision_file
temperatures = args.temperatures
plot_precision(precision_file, teffs=temperatures)
# plot_precision(precision_file, teffs=temperatures, logg=4, fe_h=1)
cumulative_plot(precision_file, teffs=temperatures)
cumulative_plot(precision_file, teffs=temperatures, full_cum=True)
sys.exit(0)