redback 1.0.0__py3-none-any.whl → 1.0.2__py3-none-any.whl

redback/transient_models/extinction_models.py

@@ -1,6 +1,5 @@
 from inspect import isfunction
 import numpy as np
-
 import redback.utils
 from redback.transient_models.fireball_models import predeceleration
 from redback.utils import logger, calc_ABmag_from_flux_density, citation_wrapper, lambda_to_nu
@@ -12,7 +11,14 @@ extinction_afterglow_base_models = ['tophat', 'cocoon', 'gaussian',
                                     'kn_afterglow', 'cone_afterglow',
                                     'gaussiancore', 'gaussian',
                                     'smoothpowerlaw', 'powerlawcore',
-                                    'tophat']
+                                    'tophat','tophat_from_emulator',
+                                    'kilonova_afterglow_redback', 'kilonova_afterglow_nakarpiran',
+                                    'tophat_redback', 'gaussian_redback', 'twocomponent_redback',
+                                    'powerlaw_redback', 'alternativepowerlaw_redback', 'doublegaussian_redback',
+                                    'tophat_redback_refreshed', 'gaussian_redback_refreshed',
+                                    'twocomponent_redback_refreshed','powerlaw_redback_refreshed',
+                                    'alternativepowerlaw_redback_refreshed', 'doublegaussian_redback_refreshed']
+
 extinction_integrated_flux_afterglow_models = extinction_afterglow_base_models
 extinction_supernova_base_models = ['sn_exponential_powerlaw', 'arnett', 'shock_cooling_and_arnett',
                                     'basic_magnetar_powered', 'slsn', 'magnetar_nickel',
@@ -20,16 +26,17 @@ extinction_supernova_base_models = ['sn_exponential_powerlaw', 'arnett', 'shock_
                                     'general_magnetar_slsn','general_magnetar_driven_supernova']
 extinction_kilonova_base_models = ['nicholl_bns', 'mosfit_rprocess', 'mosfit_kilonova',
                                    'power_law_stratified_kilonova','bulla_bns_kilonova',
-                                   'bulla_nsbh_kilonova', 'kasen_nsbh_kilonova','two_layer_stratified_kilonova',
+                                   'bulla_nsbh_kilonova', 'kasen_bns_kilonova','two_layer_stratified_kilonova',
                                    'three_component_kilonova_model', 'two_component_kilonova_model',
                                    'one_component_kilonova_model', 'one_component_ejecta_relation',
                                    'one_component_ejecta_relation_projection', 'two_component_bns_ejecta_relation',
                                    'polytrope_eos_two_component_bns', 'one_component_nsbh_ejecta_relation',
                                    'two_component_nsbh_ejecta_relation','metzger_kilonova_model']
-extinction_tde_base_models = ['tde_analytical', 'tde_semianalytical', 'gaussianrise_metzger_tde']
+extinction_tde_base_models = ['tde_analytical', 'tde_semianalytical', 'gaussianrise_cooling_envelope',
+                              'cooling_envelope', 'bpl_cooling_envelope']
 extinction_magnetar_driven_base_models = ['basic_mergernova', 'general_mergernova', 'general_mergernova_thermalisation',
                                           'general_mergernova_evolution', 'metzger_magnetar_driven_kilonova_model',
-                                          'general_metzger_magnetar_driven', 'general_magnetar_driven_thermalisation',
+                                          'general_metzger_magnetar_driven', 'general_metzger_magnetar_driven_thermalisation',
                                           'general_metzger_magnetar_driven_evolution']
 extinction_shock_powered_base_models = ['shocked_cocoon', 'shock_cooling']
 

redback/transient_models/kilonova_models.py

@@ -10,7 +10,7 @@ from redback.photosphere import TemperatureFloor, CocoonPhotosphere
 from redback.interaction_processes import Diffusion, AsphericalDiffusion
 
 from redback.utils import calc_kcorrected_properties, interpolated_barnes_and_kasen_thermalisation_efficiency, \
-    electron_fraction_from_kappa, citation_wrapper, lambda_to_nu
+    electron_fraction_from_kappa, citation_wrapper, lambda_to_nu, get_heating_terms, kappa_from_electron_fraction
 from redback.eos import PiecewisePolytrope
 from redback.sed import blackbody_to_flux_density, get_correct_output_format_from_spectra, Blackbody
 from redback.constants import *
@@ -476,7 +476,7 @@ def mosfit_rprocess(time, redshift, mej, vej, kappa, kappa_gamma, temperature_fl
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
     dense_resolution = kwargs.get('dense_resolution', 300)
-    time_temp = np.geomspace(1e-2, 5e6, dense_resolution) # in source frame
+    time_temp = np.geomspace(1e-2, 7e6, dense_resolution) # in source frame
     time_obs = time
     lbols = _mosfit_kilonova_one_component_lbol(time=time_temp,
                                                 mej=mej, vej=vej)
@@ -555,7 +555,7 @@ def mosfit_kilonova(time, redshift, mej_1, vej_1, temperature_floor_1, kappa_1,
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
     dense_resolution = kwargs.get('dense_resolution', 300)
-    time_temp = np.geomspace(1e-2, 5e6, dense_resolution)  # in source frame
+    time_temp = np.geomspace(1e-2, 7e6, dense_resolution)  # in source frame
     time_obs = time
     mej = [mej_1, mej_2, mej_3]
     vej = [vej_1, vej_2, vej_3]
@@ -817,7 +817,7 @@ def kasen_bns_kilonova(time, redshift, mej, vej, chi, **kwargs):
         fmjy = spectra.to(uu.mJy, equivalencies=uu.spectral_density(wav=output.lambdas * uu.Angstrom)).value
         nu_array = lambda_to_nu(output.lambdas)
         fmjy_func = RegularGridInterpolator((np.unique(time), nu_array), fmjy, bounds_error=False)
-        if type(frequency) == float:
+        if type(frequency) == float or type(frequency) == np.float64:
             frequency = np.ones(len(time)) * frequency
         points = np.array([time, frequency]).T
         return fmjy_func(points)
@@ -890,6 +890,11 @@ def _kilonova_hr(time, redshift, mej, velocity_array, kappa_array, beta, **kwarg
         time = time * day_to_s
         # convert to source frame time and frequency
         frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        if (isinstance(frequency, (float, int)) == False):
+            radio_mask = frequency < 10e10
+            frequency[radio_mask]=10e50
+        elif frequency < 10e10:
+            frequency =10e50
 
         _, temperature, r_photosphere = _kilonova_hr_sourceframe(time, mej, velocity_array, kappa_array, beta)
 
@@ -974,7 +979,7 @@ def three_component_kilonova_model(time, redshift, mej_1, vej_1, temperature_flo
     """
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
-    time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame
+    time_temp = np.geomspace(1e-2, 7e6, 300) # in source frame
     time_obs = time
 
     mej = [mej_1, mej_2, mej_3]
@@ -1063,7 +1068,7 @@ def two_component_kilonova_model(time, redshift, mej_1, vej_1, temperature_floor
     """
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
-    time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame
+    time_temp = np.geomspace(1e-2, 7e6, 300) # in source frame
     time_obs = time
 
     mej = [mej_1, mej_2]
@@ -1284,7 +1289,7 @@ def polytrope_eos_two_component_bns(time, redshift, mass_1, mass_2,  log_p, gamm
 
 @citation_wrapper('redback')
 def one_component_nsbh_ejecta_relation(time, redshift, mass_bh, mass_ns,
-                                        chi_eff, lambda_ns, kappa, **kwargs):
+                                        chi_bh, lambda_ns, kappa, **kwargs):
     """
     One component NSBH model
 
@@ -1292,7 +1297,7 @@ def one_component_nsbh_ejecta_relation(time, redshift, mass_bh, mass_ns,
     :param redshift: redshift
     :param mass_bh: mass of black hole
     :param mass_2: mass of neutron star
-    :param chi_eff: effective spin of black hole
+    :param chi_bh: spin of black hole along z axis
     :param lambda_ns: tidal deformability of neutron star
     :param kappa: opacity
     :param kwargs: additional keyword arguments
@@ -1308,7 +1313,7 @@ def one_component_nsbh_ejecta_relation(time, redshift, mass_bh, mass_ns,
     :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
     """
     ejecta_relation = kwargs.get('ejecta_relation', ejr.OneComponentNSBH)
-    ejecta_relation = ejecta_relation(mass_bh=mass_bh, mass_ns=mass_ns, chi_eff=chi_eff, lambda_ns=lambda_ns)
+    ejecta_relation = ejecta_relation(mass_bh=mass_bh, mass_ns=mass_ns, chi_bh=chi_bh, lambda_ns=lambda_ns)
     mej = ejecta_relation.ejecta_mass
     vej = ejecta_relation.ejecta_velocity
     output = one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs)
@@ -1316,7 +1321,7 @@ def one_component_nsbh_ejecta_relation(time, redshift, mass_bh, mass_ns,
 
 @citation_wrapper('redback')
 def two_component_nsbh_ejecta_relation(time, redshift,  mass_bh, mass_ns,
-                                        chi_eff, lambda_ns, zeta, vej_2, kappa_1, kappa_2, tf_1, tf_2, **kwargs):
+                                        chi_bh, lambda_ns, zeta, vej_2, kappa_1, kappa_2, tf_1, tf_2, **kwargs):
     """
     Two component NSBH model with dynamical and disk wind ejecta
 
@@ -1324,7 +1329,7 @@ def two_component_nsbh_ejecta_relation(time, redshift,  mass_bh, mass_ns,
     :param redshift: redshift
     :param mass_bh: mass of black hole
     :param mass_2: mass of neutron star
-    :param chi_eff: effective spin of black hole
+    :param chi_bh: spin of black hole along z axis
     :param lambda_ns: tidal deformability of neutron star
     :param zeta: fraction of disk that gets unbound
     :param vej_2: disk wind velocity in c
@@ -1344,7 +1349,7 @@ def two_component_nsbh_ejecta_relation(time, redshift,  mass_bh, mass_ns,
     :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source'
     """
     ejecta_relation = kwargs.get('ejecta_relation', ejr.TwoComponentNSBH)
-    ejecta_relation = ejecta_relation(mass_bh=mass_bh, mass_ns=mass_ns, chi_eff=chi_eff,
+    ejecta_relation = ejecta_relation(mass_bh=mass_bh, mass_ns=mass_ns, chi_bh=chi_bh,
                                       lambda_ns=lambda_ns, zeta=zeta)
     mej_1 = ejecta_relation.dynamical_mej
     mej_2 = ejecta_relation.disk_wind_mej
@@ -1376,7 +1381,7 @@ def one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs):
     """
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
-    time_temp = np.geomspace(1e-3, 5e6, 300) # in source frame
+    time_temp = np.geomspace(1e-3, 7e6, 300) # in source frame
     time_obs = time
     _, temperature, r_photosphere = _one_component_kilonova_model(time_temp, mej, vej, kappa, **kwargs)
 
@@ -1416,6 +1421,234 @@ def one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs):
                                                           spectra=spectra, lambda_array=lambda_observer_frame,
                                                           **kwargs)
 
+
+def _calc_new_heating_rate(time, mej, electron_fraction, ejecta_velocity, **kwargs):
+    """
+    Heating rate prescription following Rosswog and Korobkin 2022
+
+    :param time: time in seconds
+    :param mej: ejecta mass in solar masses
+    :param electron_fraction: electron fraction
+    :param ejecta_velocity: ejecta velocity in c
+    :param kwargs: Additional keyword arguments
+    :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate. Default is 1.0 i.e., no perturbation.
+    :param heating_rate_fudge: A fudge factor for each of the terms in the heating rate. Default to 1. i.e., no uncertainty
+    Default is 1.0 i.e., no perturbation.
+    :return: heating rate in erg/s
+    """
+    heating_terms = get_heating_terms(electron_fraction, ejecta_velocity, **kwargs)
+    heating_rate_perturbation = kwargs.get('heating_rate_perturbation', 1.0)
+    # rescale
+    m0 = mej * solar_mass
+    c1 = np.exp(heating_terms.c1)
+    c2 = np.exp(heating_terms.c2)
+    c3 = np.exp(heating_terms.c3)
+    tau1 = 1e3*heating_terms.tau1
+    tau2 = 1e5*heating_terms.tau2
+    tau3 = 1e5*heating_terms.tau3
+    term1 = 10.**(heating_terms.e0+18) * (0.5 - np.arctan((time - heating_terms.t0) / heating_terms.sig) / np.pi)**heating_terms.alp
+    term2 = (0.5 + np.arctan((time - heating_terms.t1)/heating_terms.sig1) / np.pi )**heating_terms.alp1
+    term3 = c1 * np.exp(-time/tau1)
+    term4 = c2 * np.exp(-time/tau2)
+    term5 = c3 * np.exp(-time/tau3)
+    lum_in = term1*term2 + term3 + term4 + term5
+    return lum_in*m0 * heating_rate_perturbation
+
+def _one_component_kilonova_rosswog_heatingrate(time, mej, vej, electron_fraction, **kwargs):
+    tdays = time/day_to_s
+    # set up kilonova physics
+    av, bv, dv = interpolated_barnes_and_kasen_thermalisation_efficiency(mej, vej)
+    # thermalisation from Barnes+16
+    e_th = 0.36 * (np.exp(-av * tdays) + np.log1p(2.0 * bv * tdays ** dv) / (2.0 * bv * tdays ** dv))
+    temperature_floor = kwargs.get('temperature_floor', 4000)  # kelvin
+    beta = 13.7
+
+    v0 = vej * speed_of_light
+    m0 = mej * solar_mass
+    kappa = kappa_from_electron_fraction(electron_fraction)
+    tdiff = np.sqrt(2.0 * kappa * (m0) / (beta * v0 * speed_of_light))
+
+    lum_in = _calc_new_heating_rate(time, mej, electron_fraction, vej, **kwargs)
+    integrand = lum_in * e_th * (time / tdiff) * np.exp(time ** 2 / tdiff ** 2)
+
+    bolometric_luminosity = np.zeros(len(time))
+    bolometric_luminosity[1:] = cumtrapz(integrand, time)
+    bolometric_luminosity[0] = bolometric_luminosity[1]
+    bolometric_luminosity = bolometric_luminosity * np.exp(-time ** 2 / tdiff ** 2) / tdiff
+
+    temperature = (bolometric_luminosity / (4.0 * np.pi * sigma_sb * v0 ** 2 * time ** 2)) ** 0.25
+    r_photosphere = (bolometric_luminosity / (4.0 * np.pi * sigma_sb * temperature_floor ** 4)) ** 0.5
+
+    # check temperature floor conditions
+    mask = temperature <= temperature_floor
+    temperature[mask] = temperature_floor
+    mask = np.logical_not(mask)
+    r_photosphere[mask] = v0 * time[mask]
+    return bolometric_luminosity, temperature, r_photosphere
+
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024arXiv240407271S/abstract, https://ui.adsabs.harvard.edu/abs/2024AnP...53600306R/abstract')
+def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, ye, **kwargs):
+    """
+    :param time: observer frame time in days
+    :param redshift: redshift
+    :param mej: ejecta mass in solar masses
+    :param vej: minimum initial velocity
+    :param kappa: gray opacity
+    :param kwargs: Additional keyword arguments
+    :param temperature_floor: Temperature floor in K (default 4000)
+    :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate.
+    Default is 1.0 i.e., no perturbation.
+    :param heating_rate_fudge: A fudge factor for each of the terms in the heating rate. Default to 1. i.e., no uncertainty
+    :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_temp = np.geomspace(1e-3, 7e6, 300) # in source frame
+    time_obs = time
+    _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej, vej, ye, **kwargs)
+
+    if kwargs['output_format'] == 'flux_density':
+        time = time_obs * day_to_s
+        frequency = kwargs['frequency']
+        # interpolate properties onto observation times
+        temp_func = interp1d(time_temp, y=temperature)
+        rad_func = interp1d(time_temp, y=r_photosphere)
+        # convert to source frame time and frequency
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+
+        temp = temp_func(time)
+        photosphere = rad_func(time)
+
+        flux_density = blackbody_to_flux_density(temperature=temp, r_photosphere=photosphere,
+                                                 dl=dl, frequency=frequency)
+
+        return flux_density.to(uu.mJy).value
+
+    else:
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 200))
+        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)
+        fmjy = blackbody_to_flux_density(temperature=temperature,
+                                         r_photosphere=r_photosphere, frequency=frequency[:, None], dl=dl)
+        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 get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame/day_to_s,
+                                                          spectra=spectra, lambda_array=lambda_observer_frame,
+                                                          **kwargs)
+
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024arXiv240407271S/abstract, https://ui.adsabs.harvard.edu/abs/2024AnP...53600306R/abstract')
+def two_comp_kne_rosswog_heatingrate(time, redshift, mej_1, vej_1, temperature_floor_1, ye_1,
+                                 mej_2, vej_2, temperature_floor_2, ye_2, **kwargs):
+    """
+    :param time: observer frame time in days
+    :param redshift: redshift
+    :param mej_1: ejecta mass in solar masses of first component
+    :param vej_1: minimum initial velocity of first component
+    :param kappa_1: gray opacity of first component
+    :param temperature_floor_1: floor temperature of first component
+    :param mej_2: ejecta mass in solar masses of second component
+    :param vej_2: minimum initial velocity of second component
+    :param temperature_floor_2: floor temperature of second component
+    :param kappa_2: gray opacity of second component
+    :param kwargs: Additional keyword arguments
+    :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate.
+    Default is 1.0 i.e., no perturbation.
+    :param heating_rate_fudge: A fudge factor for each of the terms in the heating rate. Default to 1. i.e., no uncertainty
+    :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_temp = np.geomspace(1e-2, 7e6, 300) # in source frame
+    time_obs = time
+
+    mej = [mej_1, mej_2]
+    vej = [vej_1, vej_2]
+    temperature_floor = [temperature_floor_1, temperature_floor_2]
+    ye = [ye_1, ye_2]
+
+    if kwargs['output_format'] == 'flux_density':
+        time = time * day_to_s
+        frequency = kwargs['frequency']
+
+        # convert to source frame time and frequency
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+
+        ff = np.zeros(len(time))
+        for x in range(2):
+            temp_kwargs = {}
+            if 'heating_rate_fudge' in kwargs:
+                temp_kwargs['heating_rate_fudge'] = kwargs['heating_rate_fudge']
+            if 'heating_rate_perturbation' in kwargs:
+                temp_kwargs['heating_rate_perturbation'] = kwargs['heating_rate_perturbation']
+            temp_kwargs['temperature_floor'] = temperature_floor[x]
+            _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej[x], vej[x], ye[x],
+                                                                          **temp_kwargs)
+            # interpolate properties onto observation times
+            temp_func = interp1d(time_temp, y=temperature)
+            rad_func = interp1d(time_temp, y=r_photosphere)
+            temp = temp_func(time)
+            photosphere = rad_func(time)
+            flux_density = blackbody_to_flux_density(temperature=temp, r_photosphere=photosphere,
+                                                     dl=dl, frequency=frequency)
+            units = flux_density.unit
+            ff += flux_density.value
+
+        ff = ff * units
+        return ff.to(uu.mJy).value
+
+    else:
+        lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 200))
+        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)
+        full_spec = np.zeros((len(time), len(frequency)))
+
+        for x in range(2):
+            temp_kwargs = {}
+            if 'heating_rate_fudge' in kwargs:
+                temp_kwargs['heating_rate_fudge'] = kwargs['heating_rate_fudge']
+            if 'heating_rate_perturbation' in kwargs:
+                temp_kwargs['heating_rate_perturbation'] = kwargs['heating_rate_perturbation']
+            temp_kwargs['temperature_floor'] = temperature_floor[x]
+            _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej[x], vej[x], ye[x],
+                                                                          **temp_kwargs)
+            fmjy = blackbody_to_flux_density(temperature=temperature,
+                                             r_photosphere=r_photosphere, frequency=frequency[:, None], dl=dl)
+            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))
+            units = spectra.unit
+            full_spec += spectra.value
+
+        full_spec = full_spec * units
+        if kwargs['output_format'] == 'spectra':
+            return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                                           lambdas=lambda_observer_frame,
+                                                                           spectra=full_spec)
+        else:
+            return get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame/day_to_s,
+                                                          spectra=full_spec, lambda_array=lambda_observer_frame,
+                                                          **kwargs)
+
 def _one_component_kilonova_model(time, mej, vej, kappa, **kwargs):
     """
     :param time: source frame time in seconds
@@ -1479,7 +1712,7 @@ def metzger_kilonova_model(time, redshift, mej, vej, beta, kappa, **kwargs):
     """
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
-    time_temp = np.geomspace(1e-4, 1e7, 300) # in source frame
+    time_temp = np.geomspace(1e-4, 7e6, 300) # in source frame
     time_obs = time
     bolometric_luminosity, temperature, r_photosphere = _metzger_kilonova_model(time_temp, mej, vej, beta,
                                                                                 kappa, **kwargs)

redback/transient_models/magnetar_driven_ejecta_models.py

@@ -58,8 +58,8 @@ def _ejecta_dynamics_and_interaction(time, mej, beta, ejecta_radius, kappa, n_is
     
     ni56_lum = 6.45e43
     co56_lum = 1.45e43
-    ni56_life = 8.8  # days
-    co56_life = 111.3  # days    
+    ni56_life = 8.8*86400  # days
+    co56_life = 111.3*86400  # days    
 
     for i in range(len(time)):
         beta = np.sqrt(1 - 1 / gamma ** 2)

redback/transient_models/phenomenological_models.py

@@ -1,5 +1,18 @@
 import numpy as np
 
+def line_spectrum(wavelength, line_amp, cont_amp, x0):
+    """
+    A gaussian to add or subtract from a continuum spectrum to mimic absorption or emission lines
+
+    :param wavelength: wavelength array in whatever units
+    :param line_amp: line amplitude scale
+    :param cont_amp: Continuum amplitude scale
+    :param x0: Position of emission line
+    :return: spectrum in whatever units set by line_amp
+    """
+    spectrum = line_amp / cont_amp * np.exp(-(wavelength - x0) ** 2. / (2 * cont_amp ** 2) )
+    return spectrum
+
 def gaussian_rise(time, a_1, peak_time, sigma_t):
     """
     :param time: time array in whatver time units

redback/transient_models/supernova_models.py

@@ -21,7 +21,7 @@ homologous_expansion_models = ['exponential_powerlaw_bolometric', 'arnett_bolome
                                'type_1c_bolometric','type_1a_bolometric']
 
 @citation_wrapper('https://zenodo.org/record/6363879#.YkQn3y8RoeY')
-def sncosmo_models(time, redshift, model_kwargs, **kwargs):
+def sncosmo_models(time, redshift, model_kwargs=None, **kwargs):
     """
     A wrapper to SNCosmo models
 
@@ -33,28 +33,36 @@ def sncosmo_models(time, redshift, model_kwargs, **kwargs):
     :param sncosmo_model: String of the SNcosmo model to use.
     :param peak_time: SNe peak time in days
     :param cosmology: astropy cosmology object by default set to Planck18
-    :param peak_abs_mag: SNe peak absolute magnitude default set to -19
-    :param peak_abs_mag_band: Band corresponding to the peak abs mag limit, default to standard::b. Must be in SNCosmo
     :param mw_extinction: Boolean for whether there is MW extinction or not. Default True
-    :param magnitude_system: Mag system; default ab
     :param host_extinction: Boolean for whether there is host extinction or not. Default True
             if used adds an extra parameter ebv which must also be in kwargs; host galaxy E(B-V). Set to 0.1 by default
+    :param use_set_peak_magnitude: Boolean for whether to set the peak magnitude or not. Default False,
+        if True the following keyword arguments also apply. Else the brightness is set by the model_kwargs.
+    :param peak_abs_mag: SNe peak absolute magnitude default set to -19
+    :param peak_abs_mag_band: Band corresponding to the peak abs mag limit, default to standard::b. Must be in SNCosmo
+    :param magnitude_system: Mag system; default ab
     :return: set by output format - 'flux_density', 'magnitude', 'flux', 'sncosmo_source'
     """
     import sncosmo
-    cosmology = kwargs.get('cosmology', cosmo)
     peak_time = kwargs.get('peak_time', 0)
-    peak_abs_mag = kwargs.get('peak_abs_mag', -19)
-    peak_abs_mag_band = kwargs.get('peak_abs_mag_band', 'standard::b')
+    cosmology = kwargs.get('cosmology', cosmo)
     model_name = kwargs.get('sncosmo_model', 'salt2')
     host_extinction = kwargs.get('host_extinction', True)
     mw_extinction = kwargs.get('mw_extinction', True)
-    magsystem = kwargs.get('magnitude_system', 'ab')
+    use_set_peak_magnitude = kwargs.get('use_set_peak_magnitude', False)
 
     model = sncosmo.Model(source=model_name)
     model.set(z=redshift)
     model.set(t0=peak_time)
-    model.update(model_kwargs)
+
+    if model_kwargs == None:
+        _model_kwargs = {}
+        for x in kwargs['model_kwarg_names']:
+            _model_kwargs[x] = kwargs[x]
+    else:
+        _model_kwargs = model_kwargs
+
+    model.update(_model_kwargs)
 
     if host_extinction:
         ebv = kwargs.get('ebv', 0.1)
@@ -63,13 +71,21 @@ def sncosmo_models(time, redshift, model_kwargs, **kwargs):
     if mw_extinction:
         model.add_effect(sncosmo.F99Dust(), 'mw', 'obs')
 
-    model.set_source_peakabsmag(peak_abs_mag, band=peak_abs_mag_band, magsys=magsystem, cosmo=cosmology)
+    if use_set_peak_magnitude:
+        peak_abs_mag = kwargs.get('peak_abs_mag', -19)
+        peak_abs_mag_band = kwargs.get('peak_abs_mag_band', 'standard::b')
+        magsystem = kwargs.get('magnitude_system', 'ab')
+        model.set_source_peakabsmag(peak_abs_mag, band=peak_abs_mag_band, magsys=magsystem, cosmo=cosmology)
 
     if kwargs['output_format'] == 'flux_density':
         frequency = kwargs['frequency']
 
+        if isinstance(frequency, float):
+            frequency = np.array([frequency])
+
         if (len(frequency) != 1 or len(frequency) == len(time)):
             raise ValueError('frequency array must be of length 1 or same size as time array')
+
         unique_frequency = np.sort(np.unique(frequency))
         angstroms = nu_to_lambda(unique_frequency)
 
@@ -90,11 +106,11 @@ def sncosmo_models(time, redshift, model_kwargs, **kwargs):
 
     if kwargs['output_format'] == 'flux':
         bands = kwargs['bands']
-        magnitude = model.bandmag(phase=time, band=bands, magsys='ab')
+        magnitude = model.bandmag(time=time, band=bands, magsys='ab')
         return sed.bandpass_magnitude_to_flux(magnitude=magnitude, bands=bands)
     elif kwargs['output_format'] == 'magnitude':
         bands = kwargs['bands']
-        magnitude = model.bandmag(phase=time, band=bands, magsys='ab')
+        magnitude = model.bandmag(time=time, band=bands, magsys='ab')
         return magnitude
     elif kwargs['output_format'] == 'sncosmo_source':
         return model
@@ -874,11 +890,9 @@ def _csm_engine(time, mej, csm_mass, vej, eta, rho, kappa, r0, **kwargs):
     # scaling constant for CSM density profile
     qq = rho * r0 ** eta
     # outer CSM shell radius
-    radius_csm = ((3.0 - eta) / (4.0 * np.pi * qq) * csm_mass + r0 ** (3.0 - eta)) ** (
-            1.0 / (3.0 - eta))
+    radius_csm = ((3.0 - eta) / (4.0 * np.pi * qq) * csm_mass + r0 ** (3.0 - eta)) ** (1.0 / (3.0 - eta))
     # photosphere radius
-    r_photosphere = abs((-2.0 * (1.0 - eta) / (3.0 * kappa * qq) +
-                         radius_csm ** (1.0 - eta)) ** (1.0 / (1.0 - eta)))
+    r_photosphere = abs((-2.0 * (1.0 - eta) / (3.0 * kappa * qq) + radius_csm ** (1.0 - eta)) ** (1.0 / (1.0 - eta)))
 
     # mass of the optically thick CSM (tau > 2/3).
     mass_csm_threshold = np.abs(4.0 * np.pi * qq / (3.0 - eta) * (
@@ -904,20 +918,19 @@ def _csm_engine(time, mej, csm_mass, vej, eta, rho, kappa, r0, **kwargs):
               (3.0 - nn) * g_n)) ** (1.0 / (3.0 - nn))) ** (
                    (nn - eta) / (eta - 3.0))
 
-    mask_1 = t_FS - time * day_to_s > 0
-    mask_2 = t_RS - time * day_to_s > 0
+    mask_RS = t_RS - time * day_to_s > 0
+    mask_FS = t_FS - time * day_to_s > 0
 
-    lbol = efficiency * (2.0 * np.pi / (nn - eta) ** 3 * g_n ** ((5.0 - eta) / (nn - eta)) * qq **
-                         ((nn - 5.0) / (nn - eta)) * (nn - 3.0) ** 2 * (nn - 5.0) * Bf ** (5.0 - eta) * AA **
-                         ((5.0 - eta) / (nn - eta)) * (time * day_to_s + ti) **
-                         ((2.0 * nn + 6.0 * eta - nn * eta - 15.) /
-                          (nn - eta)) + 2.0 * np.pi * (AA * g_n / qq) **
-                         ((5.0 - nn) / (nn - eta)) * Br ** (5.0 - nn) * g_n * ((3.0 - eta) / (nn - eta)) ** 3 * (
-                                     time * day_to_s + ti) **
-                         ((2.0 * nn + 6.0 * eta - nn * eta - 15.0) / (nn - eta)))
+    lbol_FS = 2.0 * np.pi / (nn - eta) ** 3 * g_n ** ((5.0 - eta) / (nn - eta)) * qq ** ((nn - 5.0) / (nn - eta)) \
+              * (nn - 3.0) ** 2 * (nn - 5.0) * Bf ** (5.0 - eta) * AA ** ((5.0 - eta) / (nn - eta)) * (time * day_to_s + ti) \
+              ** ((2.0 * nn + 6.0 * eta - nn * eta - 15.) / (nn - eta))
 
-    lbol[~mask_1] = 0
-    lbol[~mask_2] = 0
+    lbol_RS = 2.0 * np.pi * (AA * g_n / qq) ** ((5.0 - nn) / (nn - eta)) * Br ** (5.0 - nn) * g_n * ((3.0 - eta) / (nn - eta)) \
+              ** 3 * (time * day_to_s + ti) ** ((2.0 * nn + 6.0 * eta - nn * eta - 15.0) / (nn - eta))
+    lbol_FS[~mask_FS] = 0
+    lbol_RS[~mask_RS] = 0
+
+    lbol = efficiency * (lbol_FS + lbol_RS)
 
     csm_output = namedtuple('csm_output', ['lbol', 'r_photosphere', 'mass_csm_threshold'])
     csm_output.lbol = lbol
@@ -926,7 +939,7 @@ def _csm_engine(time, mej, csm_mass, vej, eta, rho, kappa, r0, **kwargs):
     return csm_output
 
 
-@citation_wrapper('redback')
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2013ApJ...773...76C/abstract, https://ui.adsabs.harvard.edu/abs/2017ApJ...849...70V/abstract, https://ui.adsabs.harvard.edu/abs/2020RNAAS...4...16J/abstract')
 def csm_interaction_bolometric(time, mej, csm_mass, vej, eta, rho, kappa, r0, **kwargs):
     """
     :param time: time in days in source frame
@@ -951,14 +964,15 @@ def csm_interaction_bolometric(time, mej, csm_mass, vej, eta, rho, kappa, r0, **
     csm_output = _csm_engine(time=time, mej=mej, csm_mass=csm_mass, vej=vej,
                              eta=eta, rho=rho, kappa=kappa, r0=r0, **kwargs)
     lbol = csm_output.lbol
-    r_photosphere = csm_output.r_photosphere
-    mass_csm_threshold = csm_output.mass_csm_threshold
 
     if _interaction_process is not None:
         dense_resolution = kwargs.get("dense_resolution", 1000)
-        dense_times = np.linspace(0, time[-1]+100, dense_resolution)
-        dense_lbols = _csm_engine(time=dense_times, mej=mej, csm_mass=csm_mass, vej=vej,
-                             eta=eta, rho=rho, kappa=kappa, r0=r0, **kwargs).lbol
+        dense_times = np.geomspace(0.1, time[-1]+100, dense_resolution)
+        csm_output = _csm_engine(time=dense_times, mej=mej, csm_mass=csm_mass, vej=vej,
+                                 eta=eta, rho=rho, kappa=kappa, r0=r0, **kwargs)
+        dense_lbols = csm_output.lbol
+        r_photosphere = csm_output.r_photosphere
+        mass_csm_threshold = csm_output.mass_csm_threshold
         interaction_class = _interaction_process(time=time, dense_times=dense_times, luminosity=dense_lbols,
                                                 kappa=kappa, r_photosphere=r_photosphere,
                                                 mass_csm_threshold=mass_csm_threshold, csm_mass=csm_mass, **kwargs)
@@ -966,7 +980,7 @@ def csm_interaction_bolometric(time, mej, csm_mass, vej, eta, rho, kappa, r0, **
     return lbol
 
 
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2013ApJ...773...76C/abstract')
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2013ApJ...773...76C/abstract, https://ui.adsabs.harvard.edu/abs/2017ApJ...849...70V/abstract, https://ui.adsabs.harvard.edu/abs/2020RNAAS...4...16J/abstract')
 def csm_interaction(time, redshift, mej, csm_mass, vej, eta, rho, kappa, r0, **kwargs):
     """
     :param time: time in days in observer frame
@@ -1017,7 +1031,7 @@ def csm_interaction(time, redshift, mej, csm_mass, vej, eta, rho, kappa, r0, **k
     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_temp = np.linspace(0.1, 500, 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)