"""
Plotting utility functions for GUIBRUSHR atmospheric retrieval.
This module contains visualization functions extracted from general_functions.py:
- make_smooth_mono_cmap: Create a smooth monochromatic colormap.
- plot_tp_profile: Plot the pressure-temperature (PT) profile for an atmosphere.
- plot_corner_metallicity: Generate a corner plot for molecular abundance posteriors.
- opacities_contribution_plot: Plot the contribution of different opacity sources.
- plot_vmr: Plot volume mixing ratio (VMR) profiles for different chemical species groups.
"""
import os
import pickle
from pathlib import Path
from typing import Any, Dict, Union
import matplotlib.colors as mcolors
import numpy as np
from corner import corner
from matplotlib import pyplot as plt
from scipy.interpolate import splrep, splev
# Note: GUI dependency — will be refactored to dependency injection
from GUIBRUSHR.GUI.WIDGET.MyFigure import MyFigure # subprocess-lazy-import
from GUIBRUSHR.GUI.LAYOUT.GraphicsConfig import GraphicsConfig # subprocess-lazy-import
from GUIBRUSHR.core.functions.figure_pickle import dump_figure
from GUIBRUSHR.General_Constants.Classes.UserTemperatureProfile import UserTemperatureProfile
from GUIBRUSHR.General_Constants.FunctionsAndConstants.Constant_Variables import ConstantVariables
from GUIBRUSHR.General_Constants.FunctionsAndConstants.PTR_Modules import (
split_species_all_info,
get_species_scientific_name,
get_cia_aliases,
get_species_basename,
)
from petitRADTRANS.plotlib import get_species_color
from petitRADTRANS.plotlib.style import default_color
import chemcat.utils as u
[docs]
def make_smooth_mono_cmap(color: str, name: str = 'custom') -> mcolors.LinearSegmentedColormap:
"""
Create a smooth monochromatic colormap with 5 control points.
The colormap transitions from black through dark shade, saturated color,
light shade, to white, creating a smooth gradient suitable for visualization.
Parameters:
color (str): Hex color code or named color for the saturated color point
name (str, optional): Name for the colormap. Defaults to 'custom'.
Returns:
LinearSegmentedColormap: Matplotlib colormap object with 256 color levels
"""
darkening_factor = 0.4
rgb = mcolors.to_rgb(color)
dark = tuple(c * darkening_factor for c in rgb)
light = tuple(1 - (1 - c) * darkening_factor for c in rgb)
return mcolors.LinearSegmentedColormap.from_list(
name,
['black', dark, color, light, 'white'],
N=256
)
def _init_tp_axes(pressures: np.ndarray):
"""
Create a MyFigure with axes configured for a pressure-temperature plot.
Centralizes the shared axis setup (labels, title, log pressure axis,
inverted y-limits, font sizes) so the best-fit and median-sample TP
figures stay visually identical.
Parameters:
pressures (np.ndarray): Pressure grid used to set the (inverted) y-limits.
Returns:
tuple: (root, ax) where root is the MyFigure and ax its main axis.
"""
root = MyFigure("Pressure-Temperature", type_data="Plot")
ax = root.return_ax(["PT profile"], ["Temperature [K]"], ["Pressure [bar]"])
ax.set_xlabel("Temperature [K]", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
ax.set_ylabel("Pressure [bar]", fontsize=GraphicsConfig.PLOT_FONT_SIZE_LABEL)
ax.set_title("PT profile", fontsize=GraphicsConfig.PLOT_FONT_SIZE_TITLE)
ax.tick_params(axis='both', labelsize=GraphicsConfig.PLOT_FONT_SIZE_TICK)
ax.set_yscale("log")
ax.set_ylim([max(pressures), min(pressures)])
return root, ax
[docs]
def plot_tp_profile(
range_min_pressures: float,
range_max_pressures: float,
n_layers: int,
parameters_tp_profile: Any,
gravity: float,
format_temperature: str,
path_imm: Union[str, Path],
err: Any,
rng: Any,
mcmc_chain: dict = None,
plot_median_sample: bool = False
) -> None:
"""
Plot the pressure-temperature (PT) profile for an atmosphere and save the figure.
Parameters:
range_min_pressures (float): Log10 of the minimum pressure (bar).
range_max_pressures (float): Log10 of the maximum pressure (bar).
n_layers (int): Number of pressure layers.
parameters_tp_profile (Any): Parameters for the temperature profile.
gravity (float): Surface gravity.
format_temperature (str): Name of the temperature profile method to use.
path_imm (Union[str, Path]): Path to save the output image.
err (Any): Error parameter for the temperature profile.
rng (Any): Random number generator or seed for the temperature profile.
mcmc_chain (dict, optional): Dictionary mapping parameter names to 1D arrays
of MCMC posterior samples for percentile-based error bands.
plot_median_sample (bool, optional): If True and an MCMC sample is available,
also save ``tp_profile_median_sample.jpg`` — identical to the base figure
(same 1σ/2σ bands) but with the central curve set to the per-pressure
median of the posterior temperature profiles instead of the best-fit
profile. Defaults to False.
Returns:
None
"""
pressures = np.logspace(
range_min_pressures, range_max_pressures, n_layers
)
root, ax = _init_tp_axes(pressures)
tp_profile_obj = UserTemperatureProfile(
pressures,
parameters_tp_profile,
gravity,
err,
rng,
mcmc_chain
)
# Normalize old PT profile names to new ones for backward compatibility
_pt_fallback = {"isot": "Isothermal", "madhu": "Madhusudhan", "guillot": "Guillot", "personalized": "DoubleIsothermal"}
normalized_format_temperature = _pt_fallback.get(format_temperature, format_temperature)
function_tp_profile = getattr(tp_profile_obj, normalized_format_temperature)
temperature, temp_layer_i = function_tp_profile()
ax.plot(temperature, pressures, alpha=0.5, color="red", label="PT profile")
if temp_layer_i is not None:
pressures_sample = np.logspace(
range_min_pressures, range_max_pressures, 1000
)
n_profiles = len(temp_layer_i[0, :])
n_samples_grid = 1000
temperature_sample = np.zeros([n_profiles, n_samples_grid])
for index in range(n_profiles):
temperature_sample[index, :] = splev(pressures_sample, splrep(pressures, temp_layer_i[:, index]))
# Percentile-based asymmetric sigma bands
sigma1_low = np.percentile(temperature_sample, 16, axis=0)
sigma1_high = np.percentile(temperature_sample, 84, axis=0)
sigma2_low = np.percentile(temperature_sample, 2.5, axis=0)
sigma2_high = np.percentile(temperature_sample, 97.5, axis=0)
ax.fill_betweenx(pressures_sample, sigma2_low, sigma2_high, alpha=0.3, color="blue", label=r"2$\sigma$")
ax.fill_betweenx(pressures_sample, sigma1_low, sigma1_high, alpha=0.3, color="purple", label=r"$\sigma$")
ax.legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
root.create_figure(str(Path(path_imm, "tp_profile.jpg")), 600)
# Optional companion figure: identical bands, but the central curve is the
# per-pressure median of the posterior temperature profiles (the same sample
# used for the sigma bands) instead of the best-fit profile. Only meaningful
# when an MCMC sample is available (temp_layer_i is not None).
if plot_median_sample and temp_layer_i is not None:
median_profile = np.median(temperature_sample, axis=0)
root_median, ax_median = _init_tp_axes(pressures)
ax_median.plot(
median_profile, pressures_sample, alpha=0.5, color="red",
label="PT profile (median sample)"
)
ax_median.fill_betweenx(pressures_sample, sigma2_low, sigma2_high, alpha=0.3, color="blue", label=r"2$\sigma$")
ax_median.fill_betweenx(pressures_sample, sigma1_low, sigma1_high, alpha=0.3, color="purple", label=r"$\sigma$")
ax_median.legend(fontsize=GraphicsConfig.PLOT_FONT_SIZE_LEGEND)
root_median.create_figure(str(Path(path_imm, "tp_profile_median_sample.jpg")), 600)
median_pkl_data = {
"range_min_pressures": range_min_pressures,
"range_max_pressures": range_max_pressures,
"n_layers": n_layers,
"parameters_tp_profile": parameters_tp_profile,
"gravity": gravity,
"format_temperature": format_temperature,
"pressures": pressures,
"pressures_sample": pressures_sample,
"median_profile": median_profile,
"temp_layer_i": temp_layer_i,
}
with open(str(Path(path_imm, "tp_profile_median_sample.pkl")), "wb") as _fm:
pickle.dump(median_pkl_data, _fm)
# Save plot data to pkl for later reconstruction
pkl_data = {
"range_min_pressures": range_min_pressures,
"range_max_pressures": range_max_pressures,
"n_layers": n_layers,
"parameters_tp_profile": parameters_tp_profile,
"gravity": gravity,
"format_temperature": format_temperature,
"pressures": pressures,
"temperature": temperature,
"temp_layer_i": temp_layer_i,
}
with open(str(Path(path_imm, "tp_profile.pkl")), "wb") as _f:
pickle.dump(pkl_data, _f)
[docs]
def opacities_contribution_plot(
path_results: Union[str, Path],
opacities_contribution: Dict[str, Any],
resolution: str,
function: str = "plot",
wl_scale: str = "log",
invert_y_axis: int = 0,
instruments_LR=None,
offsetLR_arr=None,
rp=None,
rs=None,
rad_mode: str = "Transmission",
) -> None:
"""
Plot and save the contribution of different opacity sources to the atmospheric model.
This function creates a visualization showing how different opacity sources (line species,
gas continua, clouds, etc.) contribute to the total atmospheric opacity. The plot is saved
as a high-resolution image in the specified output directory.
Parameters:
path_results (Union[str, Path]): Directory path where the output plot will be saved.
opacities_contribution (Dict[str, Any]): Dictionary containing opacity contributions from
different sources. Keys are source types (e.g., 'line_species', 'cloud_species'),
values are dictionaries mapping species names to their contribution data.
resolution (str): Resolution identifier used in the plot title and filename.
function (str, optional): Plotting function to use ('plot', 'scatter', etc.). Defaults to "plot".
wl_scale (str, optional): Scale for the wavelength axis ('log' or 'linear'). Defaults to "log".
invert_y_axis (int, optional): Whether to invert the y-axis. Defaults to False.
instruments_LR: Dictionary containing the instruments and their LR values.
offsetLR_arr: Array containing the offset LR values.
rp: Radius of the planet.
rs: Radius of the star.
rad_mode: Radiative mode.
Returns:
None
"""
# TODO
print("TODO opacities contribution for Emission")
if rad_mode == "Emission":
return
#
if rad_mode == "Transmission":
ylabel = r"Transit radius [$R_{{Jup}}$]"
else:
ylabel = r'1 + Emission lines'
# Validate input parameters
if wl_scale not in ['log', 'linear']:
print(f"Invalid wl_scale: {wl_scale}. Must be 'log' or 'linear'")
# Initialize color and style mappings for different opacity sources
opacity_sources_colors = {
'line_species': None,
'gas_continuum_contributors': None,
'cloud_species': None,
'rayleigh_species': None,
'opaque_cloud_top_pressure': 'darkgray',
'power_law': 'red',
'gray_opacity': 'gray',
'cloud_photosphere_median_optical_depth': 'gold',
'additional_absorption_opacities_function': 'darkviolet',
'additional_scattering_opacities_function': 'magenta',
'stellar_intensities': 'yellow',
'Total': 'k'
}
opacity_sources_linestyles = {
'line_species': '-',
'gas_continuum_contributors': '--',
'cloud_species': '-.',
'rayleigh_species': (0, (1, 0.66)), # densely dotted
'opaque_cloud_top_pressure': '-.',
'power_law': ':',
'gray_opacity': ':',
'cloud_photosphere_median_optical_depth': '-.',
'additional_absorption_opacities_function': ':',
'additional_scattering_opacities_function': ':',
'stellar_intensities': '-',
'Total': '-'
}
# Initialize color generator for species without predefined colors
default_species_color_index = (i for i in range(1001))
# Process species colors for each opacity source
for opacity_source, species_list in opacities_contribution.items():
if not isinstance(species_list, dict):
continue
if opacity_sources_colors[opacity_source] is None:
opacity_sources_colors[opacity_source] = {}
for species in species_list:
# Handle special case for gas continuum contributors
if opacity_source == 'gas_continuum_contributors':
_species = split_species_all_info(get_cia_aliases(species))[0]
_species = _species.split('--', 1)
if _species[0] == _species[1]:
_species = _species[0]
elif 'H2' in _species:
_species = _species[1] if _species[0] == 'H2' else _species[0]
else:
_species = _species[0]
else:
_species = species
# Assign color if not already assigned
if species not in opacity_sources_colors[opacity_source]:
species_color = get_species_color(get_species_basename(_species), implemented_only=False)
if species_color == default_color:
i = next(default_species_color_index)
if i == 1000:
print("Cannot support more than 1000 species")
species_color = f"C{i}"
opacity_sources_colors[opacity_source][species] = species_color
# Create figure and axis
root = MyFigure(
resolution + " Opacity Contribution",
type_data="Plot",
personal_size=[0.9, 0.6]
)
axs = root.return_ax(
[resolution + " Opacity Contribution"],
["Wavelength [nm]"],
[ylabel],
subplots=(1, 1),
figsize=(12, 6)
)
# Initialize tracking variables for y-axis limits
min_y = np.inf
max_y = 0
# Plot each opacity source
for key, value in opacities_contribution.items():
if value is None:
continue
if key == "Total":
wavelengths = value[0]
contribution = value[1]
# Plot total contribution
axs.plot(
wavelengths,
contribution,
label=key,
color=opacity_sources_colors[key],
linestyle=opacity_sources_linestyles[key]
)
# max_y = np.nanmax([max_y, np.nanmax(contribution)])
else:
if isinstance(value, dict):
# Plot individual species contributions
for species_name, species_data in value.items():
if species_data is None:
continue
# Format species label
_species = get_species_scientific_name(species_name)
if key == 'rayleigh_species':
_species += ' (Rayleigh)'
__species = species_name + ' (Rayleigh)'
elif key == 'gas_continuum_contributors':
__species = species_name.rsplit('-NatAbund', 1)[0]
else:
__species = species_name
if species_data is not None:
wavelengths = species_data[0]
contribution = species_data[1]
# Plot species contribution
getattr(axs, function)(
wavelengths,
contribution,
label=_species,
color=opacity_sources_colors[key][species_name],
linestyle=opacity_sources_linestyles[key],
alpha=0.5
)
# min_y = np.nanmin([min_y, np.nanmin(contribution)])
else:
# Plot non-species opacity sources
key_name = key.replace('_', ' ')
wavelengths = value[0]
contribution = value[1]
getattr(axs, function)(
wavelengths,
contribution,
label=key_name,
color=opacity_sources_colors[key],
linestyle=opacity_sources_linestyles[key],
alpha=0.5
)
# min_y = np.nanmin([min_y, np.nanmin(contribution)])
min_y = np.nanmin([min_y, np.nanmin(contribution)])
max_y = np.nanmax([max_y, np.nanmax(contribution)])
if instruments_LR:
for index_LR, instrument in enumerate(instruments_LR):
# Calculate offset for current instrument (first instrument has no offset)
offsetLR = 0 if (index_LR == 0 or offsetLR_arr is None) else offsetLR_arr[index_LR - 1]
# Extract instrument data
ldata = instrument.lr_data.ldata_LR
rdata = instrument.lr_data.rdata_LR - offsetLR # Shift data to model level
errdata = instrument.lr_data.errdata_LR
step = instrument.lr_data.step_LR
if rad_mode == "Transmission":
rdata_temp = 1 - rdata
rdata = np.sqrt(rdata_temp) * rs * ConstantVariables.RATIO_RSUN_RJUP_MEAN
errdata = rs * ConstantVariables.RATIO_RSUN_RJUP_MEAN * errdata / (2 * np.sqrt(rdata_temp)) # error propagation
# Plot observational data with error bars
axs.errorbar(
ldata, rdata, yerr=errdata, xerr=step, ls='none',
label=f"{instrument.name} + {-offsetLR:.5f}",
alpha=0.5, color="Red"
)
# Configure plot appearance
axs.legend(loc='center left', bbox_to_anchor=(1, 0.5))
delta_y = max_y - min_y
axs.set_ylim([min_y - delta_y * 0.05, max_y + delta_y * 0.05])
if wl_scale == 'log':
axs.set_xscale("log")
if invert_y_axis:
axs.invert_yaxis()
plt.subplots_adjust(right=0.5)
# Create output directory and save plot
path_imm = Path(path_results, "opacity_contribution")
path_imm.mkdir(parents=True, exist_ok=True)
try:
root.create_figure(str(Path(path_imm, resolution + "_img.png")), 600, False)
except IOError as e:
print(f"Failed to save opacity contribution plot: {str(e)}")
[docs]
def plot_vmr(
pressures: np.ndarray,
vmr_dict: Dict[str, np.ndarray],
folder_plot: Union[str, Path]
) -> None:
"""
Plot volume mixing ratio (VMR) profiles for different chemical species groups.
This function creates a 2x3 grid of plots showing VMR profiles for:
- All species
- HCNO neutrals
- Ions
- Alkali metals
- Metals
- Metal oxides
Parameters:
pressures (np.ndarray): Pressure values (in bar)
vmr_dict (Dict[str, np.ndarray]): Dictionary mapping species names to their VMR profiles
folder_plot (Union[str, Path]): Directory where the plot will be saved
"""
# Create figure with 2x3 subplots
root = MyFigure("VMR profiles", type_data="Plot", personal_size=[0.8, 0.8])
# Define plot labels and titles
xlabels = ["$Log_{10} VMR$"] * 6
ylabels = ["$Log_{10}$ Pressure (bar)"] * 6
titles = [
"VMR profiles all",
"VMR profiles HCNO neutrals",
"VMR profiles ions",
"VMR profiles alkali",
"VMR profiles metals",
"VMR profiles metal oxides",
]
# Create subplots
ax = root.return_ax(
titles,
xlabels,
ylabels,
subplots=[2, 3],
figsize=(14, 12),
dpi=100,
hspace=0.4,
wspace=0.3,
)
# Initialize lists for different species groups
species_groups = {
'all': [],
'HCNO_neutrals': [],
'ions': [],
'alkali': [],
'metals': [],
'metal_oxides': []
}
vmr_groups = {key: [] for key in species_groups.keys()}
# Sort species into groups
for molec, vmr in vmr_dict.items():
molec_name_compatible = molec.replace("_", "")
# Add to all species
species_groups['all'].append(molec)
vmr_groups['all'].append(vmr)
# Add to specific groups
if molec_name_compatible in ConstantVariables.HCNO_NEUTRALS:
species_groups['HCNO_neutrals'].append(molec)
vmr_groups['HCNO_neutrals'].append(vmr)
if molec_name_compatible in ConstantVariables.IONS:
species_groups['ions'].append(molec)
vmr_groups['ions'].append(vmr)
if molec_name_compatible in ConstantVariables.ALKALI:
species_groups['alkali'].append(molec)
vmr_groups['alkali'].append(vmr)
if molec_name_compatible in ConstantVariables.METALS:
species_groups['metals'].append(molec)
vmr_groups['metals'].append(vmr)
if molec_name_compatible in ConstantVariables.METAL_OXIDES:
species_groups['metal_oxides'].append(molec)
vmr_groups['metal_oxides'].append(vmr)
# Plot each group
group_positions = {
'all': (0, 0),
'HCNO_neutrals': (0, 1),
'ions': (0, 2),
'alkali': (1, 0),
'metals': (1, 1),
'metal_oxides': (1, 2)
}
for group, (row, col) in group_positions.items():
if species_groups[group]: # Only plot if there are species in the group
u.plot_vmr(
pressures,
np.transpose(vmr_groups[group]),
species_groups[group],
axis=ax[row, col],
fontsize=12
)
try:
# Save the plot
root.create_figure(str(Path(folder_plot, 'vmr_profiles.jpg')), 600)
except IOError as e:
print(f"Failed to save VMR profiles plot: {str(e)}")