"""
Module with functions for plotting the results obtained with the
:class:`~species.fit.fit_evolution.FitEvolution` class.
"""
from numbers import Real
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
from beartype import beartype
from beartype.typing import List, Optional, Tuple
from matplotlib.ticker import AutoMinorLocator
from species.read.read_isochrone import ReadIsochrone
from species.util.core_util import print_section
[docs]
@beartype
def plot_cooling(
tag: str,
n_samples: int = 50,
cooling_param: str = "log_lum",
xlim: Optional[Tuple[Real, Real]] = None,
ylim: Optional[Tuple[Real, Real]] = None,
xscale: Optional[str] = "linear",
yscale: Optional[str] = "linear",
title: Optional[str] = None,
offset: Optional[Tuple[Real, Real]] = None,
figsize: Optional[Tuple[Real, Real]] = (4.0, 2.5),
output: Optional[str] = None,
) -> Tuple[mpl.figure.Figure, List[List[List[np.ndarray]]], np.ndarray]:
"""
Function for plotting samples of cooling tracks that are
randomly drawn from the posterior distributions of the
age and mass parameters that have been estimated with
:class:`~species.fit.fit_evolution.FitEvolution`.
Parameters
----------
tag : str
Database tag where the samples are stored
n_samples : int
Number of randomly drawn cooling tracks that will be plotted.
cooling_param : str
Type of cooling parameter that will be plotted
('log_lum', 'radius', 'teff', or 'logg').
xlim : tuple(float, float), None
Limits of the wavelength axis. Automatic limits are used if
the argument is set to ``None``.
ylim : tuple(float, float), None
Limits of the flux axis. Automatic limits are used if
the argument is set to ``None``.
xscale : str, None
Scale of the x axis ('linear' or 'log'). The scale is set
to ``'linear'`` if the argument is set to ``None``.
yscale : str, None
Scale of the x axis ('linear' or 'log'). The scale is set
to ``'linear'`` if the argument is set to ``None``.
title : str
Title to show at the top of the plot.
offset : tuple(float, float)
Offset for the label of the x- and y-axis.
figsize : tuple(float, float)
Figure size.
output : str, None
Output filename for the plot. The plot is shown in an
interface window if the argument is set to ``None``.
Returns
-------
matplotlib.figure.Figure
The ``Figure`` object that can be used for further
customization of the plot.
np.ndarray
Array with the cooling tracks. The array contains
:math:`L/L_\\odot` or radius as function of time
for each companion and sample, so the shape is
(n_companions, n_samples, n_ages).
np.ndarray
Array with the random indices that have been
sampled from the posterior distribution.
"""
print_section("Plot cooling tracks")
print(f"Database tag: {tag}")
print(f"Number of samples: {n_samples}")
print(f"Model parameters: {cooling_param}")
plt.close()
if cooling_param not in ["log_lum", "luminosity", "radius", "teff", "logg"]:
raise ValueError(
"The argument of 'cooling_parameter' is "
"not valid and should be set to "
"'log_lum', 'radius', 'teff', or 'logg'."
)
from species.data.database import Database
species_db = Database()
samples_box = species_db.get_samples(tag)
samples = samples_box.samples
attr = samples_box.attributes
n_param = attr["n_param"]
n_planets = attr["n_planets"]
model_name = attr["model_name"]
log_lum = attr["log_lum"]
age_prior = attr["age_prior"]
radius_prior = attr["radius_prior"]
param_indices = {}
for i in range(n_param):
param_indices[samples_box.attributes[f"parameter{i}"]] = i
if np.isnan(age_prior[0]):
param_idx = samples_box.parameters.index("age")
age_prior = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
else:
# Asymmetric normal prior set in FitEvolution
age_prior = [
age_prior[0],
age_prior[0] + age_prior[1],
age_prior[0] + age_prior[2],
]
read_iso = ReadIsochrone(model_name)
plt.rcParams["font.family"] = "serif"
plt.rcParams["mathtext.fontset"] = "dejavuserif"
fig = plt.figure(1, figsize=figsize)
gridsp = mpl.gridspec.GridSpec(n_planets, 1)
gridsp.update(wspace=0, hspace=0.1, left=0, right=1, bottom=0, top=1)
ax = []
for i in range(n_planets):
ax.append(plt.subplot(gridsp[i, 0]))
if xscale is None:
xscale = "linear"
if yscale is None:
yscale = "linear"
cool_tracks = []
for i in range(n_planets):
cool_tracks.append([])
for i in range(n_planets):
if not isinstance(radius_prior[i], np.ndarray) and np.isnan(radius_prior[i]):
param_idx = samples_box.parameters.index(f"radius_{i}")
radius_tmp = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
else:
radius_tmp = [
radius_prior[i][0],
radius_prior[i][0] - radius_prior[i][1],
radius_prior[i][0] + radius_prior[i][1],
]
param_idx = samples_box.parameters.index(f"teff_{i}")
teff_tmp = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
param_idx = samples_box.parameters.index(f"logg_{i}")
logg_tmp = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
ax[i].set_xscale(xscale)
ax[i].set_yscale(yscale)
labelbottom = bool(i == n_planets - 1)
ax[i].tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=True,
bottom=True,
left=True,
right=True,
labelbottom=labelbottom,
)
ax[i].tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=True,
bottom=True,
left=True,
right=True,
labelbottom=labelbottom,
)
if xscale == "linear":
ax[i].xaxis.set_minor_locator(AutoMinorLocator(5))
if i == n_planets - 1:
ax[i].set_xlabel("Age (Myr)", fontsize=13)
if cooling_param in ["luminosity", "log_lum"]:
ax[i].set_ylabel(r"$\log(L/L_\odot)$", fontsize=13)
elif cooling_param == "radius":
ax[i].set_ylabel(r"Radius ($R_\mathrm{J}$)", fontsize=13)
elif cooling_param == "teff":
ax[i].set_ylabel(r"$T_\mathrm{eff}$ (K)", fontsize=13)
elif cooling_param == "logg":
ax[i].set_ylabel(r"$\log\,g$", fontsize=13)
if xlim is not None:
ax[i].set_xlim(xlim[0], xlim[1])
if ylim is not None:
ax[i].set_ylim(ylim[0], ylim[1])
if offset is not None:
ax[i].get_xaxis().set_label_coords(0.5, offset[0])
ax[i].get_yaxis().set_label_coords(offset[1], 0.5)
ran_indices = np.random.randint(low=0, high=samples.shape[0], size=n_samples)
for sample_idx in ran_indices:
for planet_idx in range(n_planets):
mass = samples[sample_idx, param_indices[f"mass_{planet_idx}"]]
if f"s_init_{planet_idx}" in param_indices:
s_init = samples[sample_idx, param_indices[f"s_init_{planet_idx}"]]
else:
s_init = None
cool_box = read_iso.get_cooling_track(
mass=mass, ages=None, s_init=s_init
)
if cooling_param in ["luminosity", "log_lum"]:
cool_tracks[planet_idx].append([cool_box.age, cool_box.log_lum])
elif cooling_param == "radius":
cool_tracks[planet_idx].append([cool_box.age, cool_box.radius])
elif cooling_param == "teff":
cool_tracks[planet_idx].append([cool_box.age, cool_box.teff])
elif cooling_param == "logg":
cool_tracks[planet_idx].append([cool_box.age, cool_box.logg])
if cool_tracks[planet_idx][-1][1] is None:
raise ValueError(
f"The selected parameter, '{cooling_param}', "
f"is not part of the '{model_name}' "
"evolutionary model grid."
)
ax[planet_idx].plot(
cool_tracks[planet_idx][-1][0],
cool_tracks[planet_idx][-1][1],
lw=0.5,
color="gray",
alpha=0.5,
)
for i in range(n_planets):
if cooling_param in ["luminosity", "log_lum"]:
ax[i].errorbar(
[age_prior[0]],
[log_lum[i][0]],
xerr=[
[age_prior[0] - np.abs(age_prior[1])],
[age_prior[2] - age_prior[0]],
],
yerr=[log_lum[i][1]],
color="tab:orange",
)
elif cooling_param == "radius":
ax[i].errorbar(
[age_prior[0]],
[radius_tmp[0]],
xerr=[
[age_prior[0] - np.abs(age_prior[1])],
[age_prior[2] - age_prior[0]],
],
yerr=[[radius_tmp[0] - radius_tmp[1]], [radius_tmp[2] - radius_tmp[0]]],
color="tab:orange",
)
elif cooling_param == "teff":
ax[i].errorbar(
[age_prior[0]],
[teff_tmp[0]],
xerr=[
[age_prior[0] - np.abs(age_prior[1])],
[age_prior[2] - age_prior[0]],
],
yerr=[[teff_tmp[0] - teff_tmp[1]], [teff_tmp[2] - teff_tmp[0]]],
color="tab:orange",
)
elif cooling_param == "logg":
ax[i].errorbar(
[age_prior[0]],
[logg_tmp[0]],
xerr=[
[age_prior[0] - np.abs(age_prior[1])],
[age_prior[2] - age_prior[0]],
],
yerr=[[logg_tmp[0] - logg_tmp[1]], [logg_tmp[2] - logg_tmp[0]]],
color="tab:orange",
)
if title is not None:
ax[0].set_title(title, fontsize=18.0)
if output is None:
plt.show()
else:
print(f"\nOutput: {output}")
plt.savefig(output, bbox_inches="tight")
return fig, cool_tracks, ran_indices
[docs]
@beartype
def plot_isochrones(
tag: str,
n_samples: int = 50,
xlim: Optional[Tuple[Real, Real]] = None,
ylim: Optional[Tuple[Real, Real]] = None,
xscale: Optional[str] = "linear",
yscale: Optional[str] = "linear",
title: Optional[str] = None,
offset: Optional[Tuple[Real, Real]] = None,
figsize: Optional[Tuple[Real, Real]] = (4.0, 2.5),
output: Optional[str] = None,
) -> Tuple[mpl.figure.Figure, List[List[List[np.ndarray]]], np.ndarray]:
"""
Function for plotting samples of isochrones that are
randomly drawn from the posterior distributions of the
age and mass parameters that have been estimated with
:class:`~species.fit.fit_evolution.FitEvolution`.
For each isochrone, the parameters from a single
posterior sample are used.
Parameters
----------
tag : str
Database tag where the samples are stored
n_samples : int
Number of randomly drawn cooling tracks that will be plotted.
xlim : tuple(float, float), None
Limits of the wavelength axis. Automatic limits are used if
the argument is set to ``None``.
ylim : tuple(float, float), None
Limits of the flux axis. Automatic limits are used if
the argument is set to ``None``.
xscale : str, None
Scale of the x axis ('linear' or 'log'). The scale is set
to ``'linear'`` if the argument is set to ``None``.
yscale : str, None
Scale of the x axis ('linear' or 'log'). The scale is set
to ``'linear'`` if the argument is set to ``None``.
title : str
Title to show at the top of the plot.
offset : tuple(float, float)
Offset for the label of the x- and y-axis.
figsize : tuple(float, float)
Figure size.
output : str, None
Output filename for the plot. The plot is shown in an
interface window if the argument is set to ``None``.
Returns
-------
matplotlib.figure.Figure
The ``Figure`` object that can be used for further
customization of the plot.
np.ndarray
Array with the isochrones. The array contains
:math:`L/L_\\odot` as function of time for each companions
and sample, so the shape is (n_companions, n_samples, n_masses).
np.ndarray
Array with the random indices that have been
sampled from the posterior distribution.
"""
print_section("Plot isochrones")
print(f"Database tag: {tag}")
print(f"Number of samples: {n_samples}")
plt.close()
from species.data.database import Database
species_db = Database()
samples_box = species_db.get_samples(tag)
samples = samples_box.samples
attr = samples_box.attributes
n_param = samples_box.attributes["n_param"]
n_planets = attr["n_planets"]
model_name = attr["model_name"]
log_lum = attr["log_lum"]
age_prior = attr["age_prior"]
mass_prior = attr["mass_prior"]
param_indices = {}
for i in range(n_param):
param_indices[samples_box.attributes[f"parameter{i}"]] = i
if np.isnan(age_prior[0]):
param_idx = samples_box.parameters.index("age")
age_prior = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
else:
# Asymmetric normal prior set in FitEvolution
age_prior = [
age_prior[0],
age_prior[0] + age_prior[1],
age_prior[0] + age_prior[2],
]
read_iso = ReadIsochrone(model_name)
plt.rcParams["font.family"] = "serif"
plt.rcParams["mathtext.fontset"] = "dejavuserif"
fig = plt.figure(1, figsize=figsize)
gridsp = mpl.gridspec.GridSpec(n_planets, 1)
gridsp.update(wspace=0, hspace=0.1, left=0, right=1, bottom=0, top=1)
ax = []
for i in range(n_planets):
ax.append(plt.subplot(gridsp[i, 0]))
if xscale is None:
xscale = "linear"
if yscale is None:
yscale = "linear"
isochrones = []
for i in range(n_planets):
isochrones.append([])
for i in range(n_planets):
labelbottom = bool(i == n_planets - 1)
ax[i].tick_params(
axis="both",
which="major",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=5,
labelsize=12,
top=True,
bottom=True,
left=True,
right=True,
labelbottom=labelbottom,
)
ax[i].tick_params(
axis="both",
which="minor",
colors="black",
labelcolor="black",
direction="in",
width=1,
length=3,
labelsize=12,
top=True,
bottom=True,
left=True,
right=True,
labelbottom=labelbottom,
)
if xscale == "linear":
ax[i].xaxis.set_minor_locator(AutoMinorLocator(5))
if i == n_planets - 1:
ax[i].set_xlabel(r"Mass ($M_\mathrm{J}$)", fontsize=13)
ax[i].set_ylabel(r"$\log(L/L_\odot)$", fontsize=13)
ax[i].set_xscale(xscale)
ax[i].set_yscale(yscale)
if xlim is not None:
ax[i].set_xlim(xlim[0], xlim[1])
if ylim is not None:
ax[i].set_ylim(ylim[0], ylim[1])
if offset is not None:
ax[i].get_xaxis().set_label_coords(0.5, offset[0])
ax[i].get_yaxis().set_label_coords(offset[1], 0.5)
ran_indices = np.random.randint(low=0, high=samples.shape[0], size=n_samples)
for sample_idx in ran_indices:
for planet_idx in range(n_planets):
age = samples[sample_idx, param_indices["age"]]
if f"s_init_{planet_idx}" in param_indices:
s_init = samples[sample_idx, param_indices[f"s_init_{planet_idx}"]]
else:
s_init = None
iso_box = read_iso.get_isochrone(
age=age, masses=None, s_init=s_init, param_interp=["log_lum"]
)
isochrones[planet_idx].append([iso_box.mass, iso_box.log_lum])
ax[planet_idx].plot(
isochrones[planet_idx][-1][0],
isochrones[planet_idx][-1][1],
lw=0.5,
color="gray",
alpha=0.5,
)
for i in range(n_planets):
if not isinstance(mass_prior[i], np.ndarray) and np.isnan(mass_prior[i]):
param_idx = samples_box.parameters.index(f"mass_{i}")
mass_tmp = np.percentile(samples[:, param_idx], [50.0, 16.0, 84.0])
else:
mass_tmp = [
mass_prior[i][0],
mass_prior[i][0] - mass_prior[i][1],
mass_prior[i][0] + mass_prior[i][1],
]
ax[i].errorbar(
[mass_tmp[0]],
[log_lum[i][0]],
xerr=[[mass_tmp[0] - mass_tmp[1]], [mass_tmp[2] - mass_tmp[0]]],
yerr=[log_lum[i][1]],
color="tab:orange",
)
if title is not None:
ax[0].set_title(title, fontsize=18.0)
if output is None:
print(f"\nOutput: {output}")
plt.show()
else:
plt.savefig(output, bbox_inches="tight")
return fig, isochrones, ran_indices
