You can interact with this notebook online: Launch notebook

This notebook is not run by nbsphinx as the data is not available but can be used as a guide.

from pathlib import Path

import numpy as np
import pandas as pd
from astropy import units as u
from matplotlib import pyplot as plt
from radioactivedecay.utils import Z_DICT
from scipy import interpolate

from tardis.io.model.readers.snec.snec_input import read_snec_isotope_profile
from tardis.io.model.readers.snec.snec_output import read_snec_output

SNEC_FOLDER_PATH = Path(
    "~/Downloads/tardis-data/MESA_STIR_MESA_SNEC"
)

Convert SNEC to TARDIS configs

input_isotope_fraction = read_snec_isotope_profile(
    SNEC_FOLDER_PATH / "input" / "profile8.data.iso.dat"
)

input_isotope_fraction_da = input_isotope_fraction.to_xr_array()

snec_output = read_snec_output(SNEC_FOLDER_PATH)

# Build the full dataset using the SNECOutput helper
snec_output_ds = snec_output.to_xr_dataset()

snec_output_ds.enclosed_mass

enclosed_mass_snec = snec_output_ds.enclosed_mass.values
interpolated_isotope_fraction = interpolate.interp1d(
    input_isotope_fraction.enclosed_mass,
    input_isotope_fraction.isotope_mass_fraction.values.T,
    bounds_error=False,
    fill_value=np.nan,
)(enclosed_mass_snec)

interpolated_isotope_fraction_df = pd.DataFrame(
    data=interpolated_isotope_fraction,
    index=input_isotope_fraction.isotope_mass_fraction.columns,
)

# Removing Neutron Fraction
interpolated_isotope_fraction_df.drop([(0, 1)], axis=0, inplace=True)

# Normalizing the isotope fractions
interpolated_isotope_fraction_df = (
    interpolated_isotope_fraction_df / interpolated_isotope_fraction_df.sum(axis=0)
)

interpolated_isotope_fraction_df

# Drop the first time step and the innermost cell
snec_output_ds = snec_output_ds.isel(time=slice(1, None), cell_id=slice(0, None))

Explorational Plotting

%matplotlib inline

fig, ax = plt.subplots(figsize=(8, 6))
vel = snec_output_ds.vel
radius = snec_output_ds.radius
for requested_t_idx in range(1, 10):
    radius_au = (radius.isel(time=requested_t_idx).values * u.cm).to(u.AU)[1:]
    vel_kms = (vel.isel(time=requested_t_idx).values * u.cm / u.s).to(u.km / u.s)[1:]
    current_time = snec_output_ds.time.values[requested_t_idx] * u.s
    homology_kms = (radius_au / current_time).to(u.km / u.s)

    ax.plot(
        radius_au.value / radius_au.max(),
        ((vel_kms - homology_kms) / homology_kms).to(1).value,
        label=f"time={current_time.to(u.day):.2f}",
        marker="|",
    )
ax.axhline(0, color="black", linestyle="--", lw=3)
ax.set_xlabel("Radius (normalized)")
ax.set_ylabel("Fractional Difference $(v - v_{hom})/v_{hom}$")
ax.set_title("Fractional Difference: Actual vs Homology Velocity")
ax.legend()
plt.tight_layout()
plt.show()
plt.savefig("homology_reached.pdf")

%matplotlib inline

fig, ax = plt.subplots(figsize=(8, 6))
ax.plot(
    (snec_output_ds.time.values[1:] * u.s).to(u.day),
    (snec_output_ds.vel_photo.values * (u.cm / u.s)).to(u.km / u.s)[1:],
)  # km/s
ax.set_xlabel("Time [days]")
ax.set_ylabel("Photosphere Velocity [km/s]")
ax.set_title("Photosphere Velocity Evolution")

Setting photospheric phase

%matplotlib inline

GAMMA_RAY_ESCAPE_THRESHOLD = 0.2  # used in Lu et al. 2024

fig, ax = plt.subplots(figsize=(8, 6))
fraction_gamma_ray = (
    snec_output_ds.lum_observed.values[1:] - snec_output_ds.lum_photo.values[1:]
) / snec_output_ds.lum_observed.values[1:]
time_days = (snec_output_ds.time.values[1:] * u.s).to(u.day).value
ax.plot(time_days, fraction_gamma_ray, label="Gamma-ray Fraction")

ax.axhline(
    GAMMA_RAY_ESCAPE_THRESHOLD,
    color="black",
    linestyle="--",
    lw=3,
    label="Threshold (20%)",
)
ax.set_xlabel("Time [days]")
ax.set_ylabel("Fraction of Escaping Gamma-ray Luminosity")
ax.set_title("Fraction of Escaping Gamma-ray Luminosity vs Time")

PHOTOSPHERE_START_INDEX = 1
PHOTOSPHERE_END_INDEX = np.abs(fraction_gamma_ray - GAMMA_RAY_ESCAPE_THRESHOLD).argmin()

# Calculate photosphere times in days
photosphere_start_day = (
    (snec_output_ds.time.values[PHOTOSPHERE_START_INDEX] * u.s).to(u.day).value
)
photosphere_end_day = (
    (snec_output_ds.time.values[PHOTOSPHERE_END_INDEX] * u.s).to(u.day).value
)

# Shade the photosphere region with axvspan
ax.axvspan(
    photosphere_start_day,
    photosphere_end_day,
    color="orange",
    alpha=0.2,
    label="Photosphere Region",
)

# Add a textbox in the middle of the shaded region
midpoint = (photosphere_start_day + photosphere_end_day) / 2
ax.text(
    midpoint,
    GAMMA_RAY_ESCAPE_THRESHOLD + 0.05,
    "Photospheric phase",
    ha="center",
    va="bottom",
    fontsize=10,
    bbox=dict(facecolor="orange", alpha=0.2, edgecolor="none"),
)

ax.legend()
plt.tight_layout()
plt.show()
plt.savefig("plots/gamma_ray_escape_photosphere.pdf")

%matplotlib inline
# Find the index in time_days closest to requested_time
requested_time = 120 * u.min
requested_t_idx = np.abs(time_days - requested_time.to(u.day).value).argmin()
x_axis_type = "vel"

if x_axis_type != "vel":
    raise ValueError("does not work - fix the Unit Conversion")
# Get the x-axis (enclosed mass) from the second time index onwards (skip time=0)
x_axis = snec_output_ds[x_axis_type].isel(time=requested_t_idx).values

if x_axis_type == "vel":
    x_axis = (x_axis * u.cm / u.s).to(u.km / u.s).value
    x_axis_label = "Velocity [km/s]"

if x_axis_type == "mass":
    x_axis = (x_axis * u.g).to(u.Msun).value
    x_axis_label = "Enclosed Mass [Msun]"

if x_axis_type == "radius":
    x_axis = (x_axis * u.cm).to(u.Rsun).value
    x_axis_label = "Radius [Rsun]"


# Convert the multi_index of interpolated_isotope_fraction_df to isotope names
df_iso = interpolated_isotope_fraction_df.copy()
df_iso.index = [f"{Z_DICT.get(z, z)}-{mass}" for z, mass in df_iso.index]
plt.clf()
fig, ax = plt.subplots(figsize=(8, 6))
for isotope in df_iso.index:
    fraction = df_iso.loc[isotope]
    ax.plot(x_axis, fraction, label=f"{isotope}", alpha=0.5, marker="|")

photo_location_last = 1e99
# Add photosphere positions as vertical lines with time labels

no_label = True
x_axis_span = x_axis.max() - x_axis.min()
for idx in range(PHOTOSPHERE_START_INDEX + 1, PHOTOSPHERE_END_INDEX + 1):
    snec_time_slice = snec_output_ds.isel(time=idx)
    photosphere_index = int(snec_time_slice.index_photo.values)
    photosphere_location = (
        (
            snec_time_slice.sel(cell_id=photosphere_index)[x_axis_type].values
            * u.cm
            / u.s
        )
        .to(u.km / u.s)
        .value
    )
    t_day = (snec_time_slice.time.values * u.s).to(u.day)
    if np.abs(photo_location_last - photosphere_location) < (x_axis_span / 20):
        continue

    # Find the cell where velocity is closest to v_photo at time=0
    vel0_kms = (vel.isel(time=0).values[1:] * u.cm / u.s).to(u.km / u.s).value

    ax.axvline(
        photosphere_location,
        color="black",
        linestyle="--",
        label="Photosphere position" if no_label else "",
        lw=1,
    )
    no_label = False
    ax.text(
        photosphere_location,
        0.5,
        f"{t_day:.1f}",
        rotation=90,
        color="black",
        va="center",
        ha="right",
        fontsize=8,
        alpha=0.7,
        transform=ax.get_xaxis_transform(),
    )
    photo_location_last = photosphere_location

ax.set_xlabel(x_axis_label)
ax.set_ylabel("Isotope Mass Fraction")
ax.set_title(
    f"{x_axis_label} @time t={requested_time:.2f}[idx={requested_t_idx}] vs Isotope Fraction"
)
ax.legend()
plt.tight_layout()
plt.show()
# plt.semilogy()
# plt.savefig("plots/isotope_fraction_photosphere_enclosed_mass.pdf")