Afterglow + Kilonova SED model

redback/priors/afterglow_kilonova_sed.prior

@@ -0,0 +1,24 @@
+redshift = Uniform(0.01, 0.3, 'redshift', latex_label=r'$z$')
+av =  Uniform(0, 2, 'av', latex_label=r'$av$')
+thv = Sine(name='thv', maximum=np.pi/2, latex_label=r'$\\theta_{\mathrm{observer}}~(\mathrm{rad})$')
+loge0 = Uniform(44, 54, 'loge0', latex_label=r'$\log_{10}~E_{0} / {\mathrm{erg}}$')
+thc = Uniform(0.01, 0.1, 'thc', latex_label=r'$\\theta_{\mathrm{core}}~({\mathrm{rad}})$')
+logn0 = Uniform(-5, 2, 'logn0', latex_label=r'$\log_{10}~n_{\mathrm{ism}} / {\mathrm{cm}}^{-3}$')
+p = Uniform(2, 3, 'p', latex_label=r'$p$')
+logepse = Uniform(-5, 0, 'logepse', latex_label=r'$\log_{10}~\\epsilon_{e}$')
+logepsb = Uniform(-5, 0, 'logepsb', latex_label=r'$\log_{10}~\\epsilon_{B}$')
+ksin = Uniform(0., 1., 'ksin', latex_label=r'$\\xi_{N}$')
+g0 = Uniform(100,2000, 'g0', latex_label=r'$\\Gamma_{0}$')
+mej_1 = Uniform(1e-2, 0.03, 'mej_1', latex_label = r'$M_{\mathrm{ej}~1}~(M_\odot)$')
+vej_1 = Uniform(0.1, 0.4, 'vej_1', latex_label = r'$v_{\mathrm{ej}~1}~(c)$')
+kappa_1 = Uniform(1, 30, 'kappa_1', latex_label = r'$\\kappa_{1}~(\mathrm{cm}^{2}/\mathrm{g})$')
+temperature_floor_1 = LogUniform(100, 6000, 'temperature_floor_1', latex_label = r'$T_{\mathrm{floor}~1}~({\mathrm{K}})$')
+mej_2 = Uniform(1e-2, 0.03, 'mej_2', latex_label = r'$M_{\mathrm{ej}~2}~(M_\odot)$')
+vej_2 = Uniform(0.1, 0.4, 'vej_2', latex_label = r'$v_{\mathrm{ej}~2}~(c)$')
+kappa_2 = Uniform(1, 30, 'kappa_2', latex_label = r'$\\kappa_{2}~(\mathrm{cm}^{2}/\mathrm{g})$')
+temperature_floor_2 = LogUniform(100, 6000, 'temperature_floor_2', latex_label = r'$T_{\mathrm{floor}~2}~({\mathrm{K}})$')
+mej_3 = Uniform(1e-2, 0.03, 'mej_2', latex_label = r'$M_{\mathrm{ej}~2}~(M_\odot)$')
+vej_3 = Uniform(0.1, 0.7, 'vej_2', latex_label = r'$v_{\mathrm{ej}~2}~(c)$')
+kappa_3 = Uniform(1, 30, 'kappa_2', latex_label = r'$\\kappa_{2}~(\mathrm{cm}^{2}/\mathrm{g})$')
+temperature_floor_3 = LogUniform(100, 6000, 'temperature_floor_2', latex_label = r'$T_{\mathrm{floor}~3}~({\mathrm{K}})$')
+kappa_gamma = DeltaFunction(10, 'kappa_gamma', latex_label = r'$\\kappa_{\\gamma}~(\mathrm{cm}^{2}/\mathrm{g})$')

redback/transient_models/combined_models.py

@@ -1,7 +1,15 @@
 import redback.transient_models.extinction_models as em
 import redback.transient_models as tm
 from redback.utils import nu_to_lambda
+from redback.utils import lambda_to_nu
+from redback.utils import day_to_s
 from redback.utils import citation_wrapper
+from redback.sed import get_correct_output_format_from_spectra
+import astropy.units as uu
+from scipy.interpolate import RegularGridInterpolator
+from collections import namedtuple
+import numpy as np
+
 
 @citation_wrapper('https://ui.adsabs.harvard.edu/abs/2020ApJ...896..166R/abstract, https://ui.adsabs.harvard.edu/abs/2020ApJ...891..152H/abstract')
 def tophat_and_twolayerstratified(time, redshift, av, thv, loge0, thc, logn0, p, logepse,
@@ -205,4 +213,100 @@ def afterglow_and_optical(time, redshift, av, **model_kwargs):
     combined = em._perform_extinction(flux_density=combined, angstroms=angstroms, av_host=av, rv_host=r_v,
                                       redshift=redshift, **kwargs)
     return combined
-
+
+@citation_wrapper('redback, and any citations for the specific model you use')
+def afterglow_kilonova_sed(time, redshift, av, **model_kwargs):
+    """
+    function to combine the flux density signals of an afterglow and kilonova model with extinction added
+
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param av: V-band extinction from host galaxy in magnitudes
+    :param model_kwargs: kwargs shared by models including frequency, lambda_array, and r_v (extinction parameter defaults to 3.1)
+    :param afterglow_kwargs: dictionary of parameters required by the afterglow transient model specified by 'base_model'
+        and any additional keyword arguments. Refer to model documentation for details.
+    :param kilonova_kwargs: dictionary of parameters required by the kilonova transient model specified by 'base_model'
+        and any additional keyword arguments. Note the base model must correspond to the given model type. Refer to model documentation
+        for details.
+    :param lambda_array: wavelength array in Angstroms, defaults to np.geomspace(100, 60000, 150)
+    :param output_format: output format ('flux_density', 'magnitude', 'spectra'), defaults to 'flux_density'
+    :return: combined afterglow and kilonova model output in the requested format. If output_format is 'spectra', 
+        returns namedtuple with time, lambdas, and spectra. Otherwise returns array in the specified format.
+    """
+
+    from redback.model_library import all_models_dict
+
+    kilonova_kwargs = model_kwargs.get('kilonova_kwargs', {'base_model': 'mosfit_kilonova'})
+    afterglow_kwargs = model_kwargs.get('afterglow_kwargs', {'base_model': 'tophat'})
+
+    temp_kwargs = model_kwargs.copy()
+    max_time = np.maximum(time.max(), 100)
+    time_observer_frame = np.geomspace(0.1, max_time, 300)
+    lambda_observer_frame = temp_kwargs.get('lambda_array', np.geomspace(100, 60000, 200))
+    frequency = lambda_to_nu(lambda_observer_frame)
+    times_mesh, frequency_mesh = np.meshgrid(time_observer_frame, frequency)
+    temp_kwargs['frequency'] = frequency_mesh
+
+    _afterglow_kwargs = afterglow_kwargs.copy()
+    _afterglow_kwargs.update(temp_kwargs)
+    _afterglow_kwargs['output_format'] = 'flux_density'
+
+    afterglow_function = all_models_dict[_afterglow_kwargs['base_model']]
+    afterglow = afterglow_function(time=times_mesh, redshift=redshift,  **_afterglow_kwargs).T
+    fmjy = afterglow * uu.mJy
+    spectra = fmjy.to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                      equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+    afterglow = namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame,
+                                                         lambdas=lambda_observer_frame,
+                                                         spectra=spectra)
+
+    _kilonova_kwargs = kilonova_kwargs.copy()
+    _kilonova_kwargs.update(temp_kwargs)
+    _kilonova_kwargs['output_format'] = 'spectra'
+
+    kilonova_function = all_models_dict[_kilonova_kwargs['base_model']]
+    capped_times = np.where(times_mesh > 7e6/day_to_s, 7e6/day_to_s, times_mesh)
+    kilonova = kilonova_function(
+        time=capped_times, 
+        redshift=redshift, **_kilonova_kwargs)
+
+    # Interpolate kilonova spectra to match afterglow's time and lambda grid
+    interpolator = RegularGridInterpolator(
+        (kilonova.time, kilonova.lambdas),
+        kilonova.spectra,
+        bounds_error=False,
+        fill_value=0.0
+    )
+
+    # Create grid points matching afterglow's time and lambda arrays
+    time_grid, lambda_grid = np.meshgrid(afterglow.time, afterglow.lambdas, indexing='ij')
+    points = np.column_stack([time_grid.ravel(), lambda_grid.ravel()])
+
+    # Interpolate kilonova to afterglow grid
+    kilonova_interpolated = interpolator(points).reshape(afterglow.spectra.shape)
+
+    combined = namedtuple('output', ['time', 'lambdas', 'spectra'])(
+        time=afterglow.time,
+        lambdas=afterglow.lambdas,
+        spectra=kilonova_interpolated + afterglow.spectra.value
+    )
+
+    # correct for host galaxy extinction
+    rest_frame_frequency = frequency * (1 + redshift)
+    r_v = model_kwargs.get('r_v', 3.1)
+    angstroms = nu_to_lambda(rest_frame_frequency)
+    combined_mJy = (combined.spectra * uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom).to(uu.mJy,
+                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom)).value
+    combined = em._perform_extinction(flux_density=combined_mJy, angstroms=angstroms, av_host=av, rv_host=r_v,
+                                      redshift=redshift, **model_kwargs)
+    fmjy = combined * uu.mJy
+    spectra = fmjy.to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom,
+                     equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom))
+    if model_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, time_eval=time_observer_frame,
+                                                      spectra=spectra, lambda_array=lambda_observer_frame,
+                                                      **model_kwargs)