redback 1.0.3__py3-none-any.whl → 1.1__py3-none-any.whl

redback/transient_models/spectral_models.py

@@ -0,0 +1,82 @@
+import numpy
+from astropy.cosmology import Planck18 as cosmo
+import redback.constants as cc
+from redback.utils import lambda_to_nu, fnu_to_flambda
+import redback.sed as sed
+import redback.transient_models.phenomenological_models as pm
+
+def _get_blackbody_spectrum(angstrom, temperature, r_photosphere, distance):
+    """
+    :param angstrom: wavelength array in angstroms
+    :param temperature: temperature in Kelvin
+    :param r_photosphere: photosphere radius in cm
+    :param distance: distance in cm
+    :return: flux in ergs/s/cm^2/angstrom
+    """
+    frequency = lambda_to_nu(angstrom)
+    flux_density = sed.blackbody_to_flux_density(frequency=frequency,
+                                                  temperature=temperature,
+                                                  r_photosphere=r_photosphere,
+                                                  dl=distance)
+    flux_density = fnu_to_flambda(f_nu=flux_density, wavelength_A=angstrom)
+    return flux_density.value
+
+def _get_powerlaw_spectrum(angstrom, alpha, aa):
+    """
+    :param angstrom: wavelength array in angstroms
+    :param alpha: power law index
+    :param aa: normalization
+    :return: flux in ergs/s/cm^2/angstrom
+    """
+    return aa*angstrom**alpha
+
+def powerlaw_spectrum_with_absorption_and_emission_lines(angstroms, alpha, aa, lc1, ls1,
+                                                         v1, lc2, ls2, v2, **kwargs):
+    """
+    A power law spectrum with one absorption line and one emission line.
+    One can add more lines if needed. Or turn the line strength to zero to remove the line.
+
+    :param angstroms: wavelength array in angstroms
+    :param alpha: power law index
+    :param aa: normalization
+    :param lc1: center of emission line
+    :param ls1: strength of emission line
+    :param v1: velocity of emission line
+    :param lc2: center of absorption line
+    :param ls2: strength of absorption line
+    :param v2: velocity of absorption line
+    :return: flux in ergs/s/cm^2/angstrom
+    """
+    flux = _get_powerlaw_spectrum(angstrom=angstroms, alpha=alpha, aa=aa)
+    fp1 = pm.line_spectrum_with_velocity_dispersion(angstroms, lc1, ls1, v1)
+    fp2 = pm.line_spectrum_with_velocity_dispersion(angstroms, lc2, ls2, v2)
+    return flux + fp1 - fp2
+
+
+def blackbody_spectrum_with_absorption_and_emission_lines(angstroms, redshift,
+                                                          rph, temp,
+                                                          lc1, ls1, v1,
+                                                          lc2, ls2, v2, **kwargs):
+    """
+    A blackbody spectrum with one absorption line and one emission line.
+    One can add more lines if needed. Or turn the line strength to zero to remove the line.
+
+    :param angstroms: wavelength array in angstroms
+    :param redshift: redshift
+    :param rph: photosphere radius in cm
+    :param temp: photosphere temperature in Kelvin
+    :param lc1: center of emission line
+    :param ls1: strength of emission line
+    :param v1: velocity of emission line
+    :param lc2: center of absorption line
+    :param ls2: strength of absorption line
+    :param v2: velocity of absorption line
+    :return: flux in ergs/s/cm^2/angstrom
+    """
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    flux = _get_blackbody_spectrum(angstrom=angstroms, distance=dl,
+                                  r_photosphere=rph, temperature=temp)
+    fp1 = pm.line_spectrum_with_velocity_dispersion(angstroms, lc1, ls1, v1)
+    fp2 = pm.line_spectrum_with_velocity_dispersion(angstroms, lc2, ls2, v2)
+    return flux + fp1 - fp2

redback/transient_models/supernova_models.py

@@ -3,7 +3,7 @@ import pandas as pd
 from redback.transient_models.phenomenological_models import exponential_powerlaw, fallback_lbol
 from redback.transient_models.magnetar_models import magnetar_only, basic_magnetar
 from redback.transient_models.magnetar_driven_ejecta_models import _ejecta_dynamics_and_interaction
-from redback.transient_models.shock_powered_models import _shock_cooling
+from redback.transient_models.shock_powered_models import _shock_cooling, _shocked_cocoon, _csm_shock_breakout
 import redback.interaction_processes as ip
 import redback.sed as sed
 import redback.photosphere as photosphere
@@ -363,7 +363,7 @@ def arnett(time, redshift, f_nickel, mej, **kwargs):
                                                               spectra=spectra, lambda_array=lambda_observer_frame,
                                                               **kwargs)
 
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract, Piro+2021')
 def shock_cooling_and_arnett(time, redshift, log10_mass, log10_radius, log10_energy,
                              f_nickel, mej, **kwargs):
     """
@@ -409,10 +409,11 @@ def shock_cooling_and_arnett(time, redshift, log10_mass, log10_radius, log10_ene
         lbol = lbol_1 + lbol_2
 
         if kwargs['interaction_process'] is not None:
-            dense_resolution = kwargs.get("dense_resolution", 1000)
-            dense_times = np.linspace(0, time[-1]+100, dense_resolution)
+            dense_resolution = kwargs.get("dense_resolution", 300)
+            dense_times = np.linspace(0.01, time[-1]+100, dense_resolution)
             dense_lbols = _nickelcobalt_engine(time=dense_times, f_nickel=f_nickel, mej=mej)
-            dense_lbols += _shock_cooling(dense_times * day_to_s, mass=mass, radius=radius, energy=energy, **kwargs).lbol
+            shock_cooling = _shock_cooling(dense_times * day_to_s, mass=mass, radius=radius, energy=energy, **kwargs).lbol
+            dense_lbols += shock_cooling
             interaction_class = kwargs['interaction_process'](time=time, dense_times=dense_times, luminosity=dense_lbols,
                                                               mej=mej, **kwargs)
             lbol = interaction_class.new_luminosity
@@ -447,6 +448,181 @@ def shock_cooling_and_arnett(time, redshift, log10_mass, log10_radius, log10_ene
                                                               spectra=spectra, lambda_array=lambda_observer_frame,
                                                               **kwargs)
 
+@citation_wrapper('https://academic.oup.com/mnras/article/522/2/2764/7086123#443111844, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+def shockcooling_morag_and_arnett_bolometric(time, v_shock, m_env, mej, f_rho, f_nickel, radius, kappa, **kwargs):
+    """
+    Assumes Shock cooling following Morag+ and arnett model for radioactive decay
+
+    :param time: time in source frame in days
+    :param v_shock: shock speed in km/s, also the ejecta velocity in the arnett calculation
+    :param m_env: envelope mass in solar masses
+    :param mej: ejecta mass in solar masses
+    :param f_rho: f_rho. Typically, of order unity
+    :param f_nickel: fraction of nickel mass
+    :param radius: star/envelope radius in units of 10^13 cm
+    :param kappa: opacity in cm^2/g
+    :param kwargs: Additional parameters required by model
+    :return: bolometric luminosity in erg/s
+    """
+    from redback.transient_models.shock_powered_models import shockcooling_morag_bolometric
+    f_rho_m = f_rho * mej
+    nickel_lbol = arnett_bolometric(time=time, f_nickel=f_nickel,
+                                    mej=mej, interaction_process=ip.Diffusion, kappa=kappa, vej=v_shock, **kwargs)
+    sbo_output = shockcooling_morag_bolometric(time=time, v_shock=v_shock, m_env=m_env, f_rho_m=f_rho_m,
+                                                     radius=radius, kappa=kappa, **kwargs)
+    lbol = nickel_lbol + sbo_output
+    return lbol
+
+@citation_wrapper('https://academic.oup.com/mnras/article/522/2/2764/7086123#443111844, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+def shockcooling_morag_and_arnett(time, redshift, v_shock, m_env, mej, f_rho, f_nickel, radius, kappa, **kwargs):
+    """
+    Assumes Shock cooling following Morag+ and arnett model for radioactive decay
+
+    :param time: time in observer frame in days
+    :param redshift: source redshift
+    :param v_shock: shock speed in km/s, also the ejecta velocity in the arnett calculation
+    :param m_env: envelope mass in solar masses
+    :param mej: ejecta mass in solar masses
+    :param f_rho: f_rho. Typically, of order unity
+    :param f_nickel: fraction of nickel mass
+    :param radius: star/envelope radius in units of 10^13 cm
+    :param kappa: opacity in cm^2/g
+    :param kwargs: Additional parameters required by model
+    :param frequency: Required if output_format is 'flux_density'.
+        frequency to calculate - Must be same length as time array or a single number).
+    :param bands: Required if output_format is 'magnitude' or 'flux'.
+    :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on.
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    """
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        lbol = shockcooling_morag_and_arnett_bolometric(time=time, v_shock=v_shock, m_env=m_env, mej=mej,
+                                                        f_rho=f_rho, f_nickel=f_nickel, radius=radius, kappa=kappa, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_shock, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature, r_photosphere=photo.r_photosphere,
+                              frequency=frequency, luminosity_distance=dl)
+        flux_density = sed_1.flux_density
+        return flux_density.to(uu.mJy).value
+    else:
+        time_obs = time
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 100))
+        time_temp = np.geomspace(0.1, 3000, 300)  # in days
+        time_observer_frame = time_temp * (1. + redshift)
+        frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
+                                                     redshift=redshift, time=time_observer_frame)
+        lbol = shockcooling_morag_and_arnett_bolometric(time=time, v_shock=v_shock, m_env=m_env, mej=mej,
+                                                        f_rho=f_rho, f_nickel=f_nickel, radius=radius, kappa=kappa, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_shock, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature, r_photosphere=photo.r_photosphere,
+                              frequency=frequency[:, None], luminosity_distance=dl)
+        fmjy = sed_1.flux_density.T
+        spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                        lambdas=lambda_observer_frame,
+                                                                        spectra=spectra)
+        else:
+            return sed.get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame,
+                                                              spectra=spectra, lambda_array=lambda_observer_frame,
+                                                              **kwargs)
+
+@citation_wrapper('https://iopscience.iop.org/article/10.3847/1538-4357/aa64df, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+def shockcooling_sapirwaxman_and_arnett_bolometric(time, v_shock, m_env, mej, f_rho, f_nickel, radius, kappa, **kwargs):
+    """
+    Assumes Shock cooling following Sapir and Waxman and arnett model for radioactive decay
+
+    :param time: time in source frame in days
+    :param v_shock: shock speed in km/s, also the ejecta velocity in the arnett calculation
+    :param m_env: envelope mass in solar masses
+    :param mej: ejecta mass in solar masses
+    :param f_rho: f_rho. Typically, of order unity
+    :param f_nickel: fraction of nickel mass
+    :param radius: star/envelope radius in units of 10^13 cm
+    :param kappa: opacity in cm^2/g
+    :param kwargs: Additional parameters required by model
+    :param n: index of progenitor density profile, 1.5 (default) or 3.0
+    :param RW: If True, use the simplified Rabinak & Waxman formulation (off by default)
+    :return: bolometric luminosity in erg/s
+    """
+    from redback.transient_models.shock_powered_models import shockcooling_sapirandwaxman_bolometric
+    f_rho_m = f_rho * mej
+    nickel_lbol = arnett_bolometric(time=time, f_nickel=f_nickel,
+                                    mej=mej, interaction_process=ip.Diffusion, kappa=kappa, vej=v_shock, **kwargs)
+    sbo_output = shockcooling_sapirandwaxman_bolometric(time=time, v_shock=v_shock, m_env=m_env, f_rho_m=f_rho_m,
+                                                     radius=radius, kappa=kappa, **kwargs)
+    lbol = nickel_lbol + sbo_output
+    return lbol
+
+@citation_wrapper('https://iopscience.iop.org/article/10.3847/1538-4357/aa64df, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+def shockcooling_sapirwaxman_and_arnett(time, redshift, v_shock, m_env, mej, f_rho, f_nickel, radius, kappa, **kwargs):
+    """
+    Assumes Shock cooling following Sapir and Waxman and arnett model for radioactive decay
+
+    :param time: time in source frame in days
+    :param v_shock: shock speed in km/s, also the ejecta velocity in the arnett calculation
+    :param m_env: envelope mass in solar masses
+    :param mej: ejecta mass in solar masses
+    :param f_rho: f_rho. Typically, of order unity
+    :param f_nickel: fraction of nickel mass
+    :param radius: star/envelope radius in units of 10^13 cm
+    :param kappa: opacity in cm^2/g
+    :param kwargs: Additional parameters required by model
+    :param n: index of progenitor density profile, 1.5 (default) or 3.0
+    :param RW: If True, use the simplified Rabinak & Waxman formulation (off by default)
+    :param frequency: Required if output_format is 'flux_density'.
+        frequency to calculate - Must be same length as time array or a single number).
+    :param bands: Required if output_format is 'magnitude' or 'flux'.
+    :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on.
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    """
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        lbol = shockcooling_morag_and_arnett_bolometric(time=time, v_shock=v_shock, m_env=m_env, mej=mej,
+                                                        f_rho=f_rho, f_nickel=f_nickel, radius=radius, kappa=kappa,
+                                                        **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_shock, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature, r_photosphere=photo.r_photosphere,
+                              frequency=frequency, luminosity_distance=dl)
+        flux_density = sed_1.flux_density
+        return flux_density.to(uu.mJy).value
+    else:
+        time_obs = time
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 100))
+        time_temp = np.geomspace(0.1, 3000, 300)  # in days
+        time_observer_frame = time_temp * (1. + redshift)
+        frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
+                                                     redshift=redshift, time=time_observer_frame)
+        lbol = shockcooling_morag_and_arnett_bolometric(time=time, v_shock=v_shock, m_env=m_env, mej=mej,
+                                                        f_rho=f_rho, f_nickel=f_nickel, radius=radius, kappa=kappa,
+                                                        **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_shock, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature, r_photosphere=photo.r_photosphere,
+                              frequency=frequency[:, None], luminosity_distance=dl)
+        fmjy = sed_1.flux_density.T
+        spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                        lambdas=lambda_observer_frame,
+                                                                        spectra=spectra)
+        else:
+            return sed.get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame,
+                                                              spectra=spectra, lambda_array=lambda_observer_frame,
+                                                              **kwargs)
+
 @citation_wrapper('redback')
 def basic_magnetar_powered_bolometric(time, p0, bp, mass_ns, theta_pb, **kwargs):
     """
@@ -707,7 +883,6 @@ def magnetar_nickel(time, redshift, f_nickel, mej, p0, bp, mass_ns, theta_pb, **
                                                               **kwargs)
 
 
-
 @citation_wrapper('redback')
 def homologous_expansion_supernova_model_bolometric(time, mej, ek, **kwargs):
     """
@@ -1619,12 +1794,12 @@ def csm_shock_and_arnett_bolometric(time, mej, f_nickel, csm_mass, v_min, beta,
     lbol = nickel_lbol + sbo_output
     return lbol
 
-
 @citation_wrapper('https://ui.adsabs.harvard.edu/abs/2022ApJ...933..238M/abstract, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
 def csm_shock_and_arnett(time, redshift, mej, f_nickel, csm_mass, v_min, beta, shell_radius,
                shell_width_ratio, kappa, **kwargs):
     """
     Assumes CSM interaction for a shell-like CSM with a hard outer boundary and arnett model for radioactive decay
+    Assumes one single photosphere from the sum of the bolometric luminosities
 
     :param time: time in days in observer frame
     :param redshift: source redshift
@@ -1674,6 +1849,157 @@ def csm_shock_and_arnett(time, redshift, mej, f_nickel, csm_mass, v_min, beta, s
         sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature, r_photosphere=photo.r_photosphere,
                               frequency=frequency[:, None], luminosity_distance=dl)
         fmjy = sed_1.flux_density.T
+        spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                        lambdas=lambda_observer_frame,
+                                                                        spectra=spectra)
+        else:
+            return sed.get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame,
+                                                              spectra=spectra, lambda_array=lambda_observer_frame,
+                                                              **kwargs)
+
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2022ApJ...933..238M/abstract, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
+def csm_shock_and_arnett_two_rphots(time, redshift, mej, f_nickel, csm_mass, v_min, beta, shell_radius,
+               shell_width_ratio, kappa, **kwargs):
+    """
+    Assumes CSM interaction for a shell-like CSM with a hard outer boundary and arnett model for radioactive decay.
+    Assumes the photospheres for the CSM-interaction and the Arnett model are different.
+
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param mej: ejecta mass in solar masses
+    :param f_nickel: fraction of nickel mass
+    :param csm_mass: csm mass in solar masses
+    :param v_min: ejecta velocity in km/s
+    :param beta: velocity ratio in c (beta < 1)
+    :param shell_radius: radius of shell in 10^14 cm
+    :param kappa: opacity
+    :param shell_width_ratio: shell width ratio (deltaR/R0)
+    :param kwargs: kappa_gamma, temperature_floor, and any kwarg to
+                change any other input physics/parameters from default.
+    :param frequency: Required if output_format is 'flux_density'.
+        frequency to calculate - Must be same length as time array or a single number).
+    :param bands: Required if output_format is 'magnitude' or 'flux'.
+    :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on.
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    """
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+
+        output = _csm_shock_breakout(time=time, csm_mass=csm_mass*solar_mass, v_min=v_min, beta=beta, kappa=kappa,
+                                     shell_radius=shell_radius, shell_width_ratio=shell_width_ratio, **kwargs)
+        r_phot = output.r_photosphere
+        temp = output.temperature
+        flux_density = sed.blackbody_to_flux_density(temperature=temp, r_photosphere=r_phot, dl=dl, frequency=frequency)
+        lbol = arnett_bolometric(time=time, f_nickel=f_nickel, mej=mej, vej=v_min, kappa=kappa,
+                                 interaction_process=ip.Diffusion, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_min, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature,
+                              r_photosphere=photo.r_photosphere, frequency=frequency, luminosity_distance=dl)
+        flux_density += sed_1.flux_density
+        return flux_density.to(uu.mJy).value
+    else:
+        time_obs = time
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 100))
+        time_temp = np.geomspace(0.1, 300, 300)  # in days
+        time_observer_frame = time_temp * (1. + redshift)
+        frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
+                                                     redshift=redshift, time=time_observer_frame)
+        output = _csm_shock_breakout(time=time, csm_mass=csm_mass * solar_mass, v_min=v_min, beta=beta, kappa=kappa,
+                                     shell_radius=shell_radius, shell_width_ratio=shell_width_ratio, **kwargs)
+        fmjy = sed.blackbody_to_flux_density(temperature=output.temperature,
+                                             r_photosphere=output.r_photosphere, frequency=frequency[:, None], dl=dl)
+        lbol = arnett_bolometric(time=time, f_nickel=f_nickel, mej=mej, vej=v_min, kappa=kappa,
+                                 interaction_process=ip.Diffusion, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=v_min, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature,
+                              r_photosphere=photo.r_photosphere, frequency=frequency[:, None], luminosity_distance=dl)
+        fmjy += sed_1.flux_density
+        fmjy = fmjy.T
+        spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                        lambdas=lambda_observer_frame,
+                                                                        spectra=spectra)
+        else:
+            return sed.get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame,
+                                                              spectra=spectra, lambda_array=lambda_observer_frame,
+                                                              **kwargs)
+
+def shocked_cocoon_and_arnett(time, redshift, mej_c, vej_c, eta, tshock, shocked_fraction, cos_theta_cocoon, kappa,
+                              mej, f_nickel, vej, **kwargs):
+    """
+    Emission from a shocked cocoon and arnett model for radioactive decay.
+    We assume two different photospheres here.
+
+    :param time: Time in days in observer frame
+    :param redshift: redshift
+    :param mej_c: cocoon mass (in solar masses)
+    :param vej_c: cocoon material velocity (in c)
+    :param eta: slope for the cocoon density profile
+    :param tshock: shock breakout time (in seconds)
+    :param shocked_fraction: fraction of the cocoon shocked
+    :param cos_theta_cocoon: cosine of the cocoon opening angle
+    :param kappa: opacity
+    :param mej: supernova ejecta mass (in solar masses)
+    :param f_nickel: fraction of nickel for ejecta mass
+    :param vej: supernova ejecta velocity (in km/s)
+    :param kwargs: Extra parameters used by model e.g., kappa_gamma, temperature_floor, and any kwarg to
+                change any other input physics/parameters from default.
+    :param frequency: Required if output_format is 'flux_density'.
+        frequency to calculate - Must be same length as time array or a single number).
+    :param bands: Required if output_format is 'magnitude' or 'flux'.
+    :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on.
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
+    """
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    time_obs = time
+
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        output = _shocked_cocoon(time=time, mej=mej_c, vej=vej_c, eta=eta,
+                                 tshock=tshock, shocked_fraction=shocked_fraction,
+                                 cos_theta_cocoon=cos_theta_cocoon, kappa=kappa)
+        flux_density = sed.blackbody_to_flux_density(temperature=output.temperature, r_photosphere=output.r_photosphere,
+                                                     dl=dl, frequency=frequency)
+        lbol = arnett_bolometric(time=time, f_nickel=f_nickel, mej=mej, vej=vej,
+                                 interaction_process=ip.Diffusion, kappa=kappa, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=vej, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature,
+                              r_photosphere=photo.r_photosphere, frequency=frequency, luminosity_distance=dl)
+        flux_density += sed_1.flux_density
+        return flux_density.to(uu.mJy).value
+    else:
+        lambda_observer_frame = kwargs.get('frequency_array', np.geomspace(100, 60000, 200))
+        time_temp = np.linspace(1e-2, 300, 300)
+        time_observer_frame = time_temp
+        frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
+                                                     redshift=redshift, time=time_observer_frame)
+        output = _shocked_cocoon(time=time, mej=mej_c, vej=vej_c, eta=eta,
+                                 tshock=tshock, shocked_fraction=shocked_fraction,
+                                 cos_theta_cocoon=cos_theta_cocoon, kappa=kappa)
+        fmjy = sed.blackbody_to_flux_density(temperature=output.temperature,
+                                             r_photosphere=output.r_photosphere, frequency=frequency[:, None], dl=dl)
+        lbol = arnett_bolometric(time=time, f_nickel=f_nickel, mej=mej, vej=vej,
+                                 interaction_process=ip.Diffusion, kappa=kappa, **kwargs)
+        photo = photosphere.TemperatureFloor(time=time, luminosity=lbol, vej=vej, **kwargs)
+        sed_1 = sed.Blackbody(temperature=photo.photosphere_temperature,
+                              r_photosphere=photo.r_photosphere, frequency=frequency[:, None], luminosity_distance=dl)
+        fmjy += sed_1.flux_density
+        fmjy = fmjy.T
         spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
                                      equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
         if kwargs['output_format'] == 'spectra':

redback/transient_models/tde_models.py

@@ -1178,8 +1178,8 @@ def fitted_pl_decay(time, redshift, log_mh, a_bh, m_disc, r0, tvi, t_form, incl,
         nulnus_decay = np.zeros(len(time))
         
         if len(freqs_un) == 1:
-            nulnus_plateau = m.model_UV(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, frequency)
-            nulnus_decay = m.decay_model(time, log_L, tdecay, p, t_peak, log_T, v=freqs_un[0])
+            nulnus_plateau = m.model_UV(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, v=freqs_un[0])
+            nulnus_decay = m.decay_model(time, log_L, t_decay, p, t_peak, log_T, v=freqs_un[0])
             nulnus_rise = m.rise_model(time, log_L, sigma, t_peak, log_T, v=freqs_un[0])
         else:
             for i in range(0,len(freqs_un)):
@@ -1199,9 +1199,11 @@ def fitted_pl_decay(time, redshift, log_mh, a_bh, m_disc, r0, tvi, t_form, incl,
         frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
                                                      redshift=redshift, time=time_observer_frame)
         nulnus_plateau = m.model_SEDs(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, frequency)
-        nulnus_risedecay = np.zeros((100, 300))
-        for i in range(0,len(frequency)):
-            nulnus_risedecay[i,:] = m.decay_model(time, log_L, t_decay, p, t_peak, log_T, v=frequency[i]) + m.rise_model(time, log_L, sigma, t_peak, log_T, v=frequency[i])  
+
+        freq_0 = 6e14
+        l_e_amp = (model.decay_model(time, log_L, t_decay, t_peak, log_T, freq_0) + model.rise_model(time, log_L, sigma, t_peak, log_T, freq_0))
+        nulnus_risedecay = ((l_e_amp[:, None] * (frequency/freq_0)**4 * 
+                        (np.exp(cc.planck * freq_0/(cc.boltzmann_constant * 10**log_T)) - 1)/(np.exp(cc.planck * frequency/(cc.boltzmann_constant * 10**log_T)) - 1)).T)  
         flux_density = ((nulnus_risedecay + nulnus_plateau)/(4.0 * np.pi * dl**2 * frequency[:,np.newaxis] * 1.0e-26))  
         fmjy = flux_density.T           
         spectra = (fmjy * uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
@@ -1259,8 +1261,8 @@ def fitted_exp_decay(time, redshift, log_mh, a_bh, m_disc, r0, tvi, t_form, incl
         nulnus_decay = np.zeros(len(time))
         
         if len(freqs_un) == 1:
-            nulnus_plateau = m.model_UV(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, frequency)
-            nulnus_decay = m.decay_model(time, log_L, tdecay, t_peak, log_T, v=freqs_un[0])
+            nulnus_plateau = m.model_UV(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, v=freqs_un[0])
+            nulnus_decay = m.decay_model(time, log_L, t_decay, t_peak, log_T, v=freqs_un[0])
             nulnus_rise = m.rise_model(time, log_L, sigma, t_peak, log_T, v=freqs_un[0])
         else:
             for i in range(0,len(freqs_un)):
@@ -1280,9 +1282,11 @@ def fitted_exp_decay(time, redshift, log_mh, a_bh, m_disc, r0, tvi, t_form, incl
         frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
                                                      redshift=redshift, time=time_observer_frame)
         nulnus_plateau = m.model_SEDs(time, log_mh, a_bh, m_disc, r0, tvi, t_form, ang, frequency)
-        nulnus_risedecay = np.zeros((100, 300))
-        for i in range(0,len(frequency)):
-            nulnus_risedecay[i,:] = m.decay_model(time, log_L, t_decay, t_peak, log_T, v=frequency[i]) + m.rise_model(time, log_L, sigma, t_peak, log_T, v=frequency[i])   
+
+        freq_0 = 6e14
+        l_e_amp = (m.decay_model(time, log_L, t_decay, t_peak, log_T, freq_0) + m.rise_model(time, log_L, sigma, t_peak, log_T, freq_0))
+        nulnus_risedecay = ((l_e_amp[:, None] * (frequency/freq_0)**4 * 
+                        (np.exp(cc.planck * freq_0/(cc.boltzmann_constant * 10**log_T)) - 1)/(np.exp(cc.planck * frequency/(cc.boltzmann_constant * 10**log_T)) - 1)).T) 
         flux_density = ((nulnus_risedecay + nulnus_plateau)/(4.0 * np.pi * dl**2 * frequency[:,np.newaxis] * 1.0e-26))  
         fmjy = flux_density.T           
         spectra = (fmjy * uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
@@ -1378,11 +1382,10 @@ def stream_stream_tde_bolometric(time, mbh_6, mstar, c1, f, h_r, inc_tcool, del_
     norm = output.bolometric_luminosity[0] / f1
 
     #evaluate giant array of bolometric luminosities
-    tt_pre_fb = np.linspace(0, output.time_temp[0], 100)
+    tt_pre_fb = np.linspace(0, output.time_temp[0]-0.001, 100)
     tt_post_fb = output.time_temp
     full_time = np.concatenate([tt_pre_fb, tt_post_fb])
-    f1 = pm.gaussian_rise(time=tt_pre_fb, a_1=norm,
-                          peak_time=peak_time * cc.day_to_s, sigma_t=sigma_t * cc.day_to_s)
+    f1 = norm * np.exp(-(tt_pre_fb - (peak_time * cc.day_to_s))**2.0 / (2 * (sigma_t * cc.day_to_s) **2.0))
     f2 = output.bolometric_luminosity
     full_lbol = np.concatenate([f1, f2])
     lbol_func = interp1d(full_time, y=full_lbol, fill_value='extrapolate')
@@ -1420,7 +1423,7 @@ def stream_stream_tde(time, redshift, mbh_6, mstar, c1, f, h_r, inc_tcool, del_o
     #get bolometric and temperature info
     f1 = pm.gaussian_rise(time=output.time_temp[0] / cc.day_to_s, a_1=1, peak_time=peak_time, sigma_t=sigma_t)
     norm = output.bolometric_luminosity[0] / f1    
-    tt_pre_fb = np.linspace(0, output.time_temp[0], 100)
+    tt_pre_fb = np.linspace(0, output.time_temp[0]-0.001, 100)
     tt_post_fb = output.time_temp
     full_time = np.concatenate([tt_pre_fb, tt_post_fb])
     f1_src = pm.gaussian_rise(time=tt_pre_fb, a_1=norm,
@@ -1436,34 +1439,47 @@ def stream_stream_tde(time, redshift, mbh_6, mstar, c1, f, h_r, inc_tcool, del_o
     if kwargs['output_format'] == 'flux_density':
         frequency = kwargs['frequency']
         if isinstance(frequency, float):
-            frequency = np.ones(len(time)) * frequency
-            
-    else:
-        bands = kwargs['bands']
-        if isinstance(bands, str):
-            bands = [str(bands) for x in range(len(time))]
-        frequency=bands_to_frequency(bands)           
+            frequency = np.ones(len(time)) * frequency           
     
-    # convert to source frame time and frequency
-    frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
-    unique_frequency = np.sort(np.unique(frequency))
-
-    # build flux density function for each frequency
-    flux_den_interp_func = {}
-    total_time = full_time * (1 + redshift)
-    for freq in unique_frequency:           
-        flux_den = sed.blackbody_to_flux_density(temperature=full_temp,
+        # convert to source frame time and frequency
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        unique_frequency = np.sort(np.unique(frequency))
+
+        # build flux density function for each frequency
+        flux_den_interp_func = {}
+        total_time = full_time * (1 + redshift)
+        for freq in unique_frequency:           
+            flux_den = sed.blackbody_to_flux_density(temperature=full_temp,
                                            r_photosphere=r_eff,
                                            dl=dl, frequency=freq).to(uu.mJy)
-        flux_den_interp_func[freq] = interp1d(total_time, flux_den, fill_value='extrapolate')
+            flux_den_interp_func[freq] = interp1d(total_time, flux_den, fill_value='extrapolate')
 
-    # interpolate onto actual observed frequency and time values
-    flux_density = []
-    for freq, tt in zip(frequency, time):
-        flux_density.append(flux_den_interp_func[freq](tt * cc.day_to_s))
-    flux_density = flux_density * uu.mJy
-    
-    if kwargs['output_format'] == 'flux_density':    
-        return flux_density.to(uu.mJy).value        
-    else:            
-        return calc_ABmag_from_flux_density(flux_density.to(uu.mJy).value).value
+        # interpolate onto actual observed frequency and time values
+        flux_density = []
+        for freq, tt in zip(frequency, time):
+            flux_density.append(flux_den_interp_func[freq](tt * cc.day_to_s))
+        flux_density = flux_density * uu.mJy
+        return flux_density.to(uu.mJy).value    
+        
+    else:
+        time_obs = time
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 100))
+        time_observer_frame = full_time * (1. + redshift)
+        frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame),
+                                                     redshift=redshift, time=time_observer_frame)
+        freq_0 = 6e14
+        flux_den = sed.blackbody_to_flux_density(temperature=full_temp,
+                            r_photosphere=r_eff,
+                            dl=dl, frequency=freq_0).to(uu.mJy)           
+        fmjy = ((flux_den[:, None] * (frequency/freq_0)**4 * 
+                        (np.exp(cc.planck * freq_0/(cc.boltzmann_constant * full_temp[:, None])) - 1) / (np.exp(cc.planck * frequency/(cc.boltzmann_constant * full_temp[:, None])) - 1)).T)                
+        spectra = fmjy.T.to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom)) 
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                          lambdas=lambda_observer_frame,
+                                                                          spectra=spectra)
+        else:
+            return sed.get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame/cc.day_to_s,
+                                                              spectra=spectra, lambda_array=lambda_observer_frame,
+                                                              **kwargs)