# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:desi] *
#     language: python
#     name: conda-env-desi-py
# ---

# # Transient Models
#
# Example of randomly grabbing transient models by type from the `desisim.transients` module.
#
# Transient models can also be accessed by name, which is demonstrated below.

# +
import numpy as np
import matplotlib.pyplot as plt
from astropy import units as u

from desisim.transients import transients


# -

# ## Plotting Script
#
# Given a transient model and a matplotlib Axes object, plot the flux vs. wavelength at peak flux, always defined to be at time 0.

def plot_transient(tr, ax):
    wlmin = np.maximum(2500., tr.minwave().value)
    wlmax = np.minimum(9500., tr.maxwave().value)
    wl = np.arange(wlmin, wlmax, 1.) * u.Angstrom
    t = 0*u.day
    ax.plot(wl, tr.flux(0*u.day, wl),
            label='{} ({})'.format(tr.model, tr.type))
    ax.set(xlabel=r'wavelength [$\AA$]',
           ylabel=r'flux [arb. units]')
    ax.legend(fontsize=10)


# ## Available Models
#
# List transient models. Then pick one and plot it.

print(transients)

fig, ax = plt.subplots(1,1, figsize=(6,4), tight_layout=True)
model = transients.get_model('hsiao')
print(model.mintime(), model.maxtime())
plot_transient(model, ax)

# ## Model Random Access
#
# You can grab a model by type. The `desisim.transients` model will randomly select one of the avaialable models of that type. Types are just strings, and if you try one that is not available in the model you'll just get an exception.
#
# Below, loop over several types, randomly select a model from each, and plot the spectrum from t0 (max light) for that model.

types = ['Ia', 'Ib', 'Ib/c', 'Ic', 'IIn', 'IIP', 'IIL', 'IIL/P', 'II-pec']
ntype = len(types)

# +
ncol = 3
nrow = int(np.ceil(ntype / ncol))
fig, axes = plt.subplots(nrow, ncol, figsize=(4*ncol, 3*nrow), tight_layout=True)
axes = axes.flatten()

for t, ax in zip(types, axes):
    s = transients.get_type(t)
    print(vars(s))
    plot_transient(s, ax)
# -



StarcoderdataJupyterScriptsDedupFiltered