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

redback/transient_models/afterglow_models.py

@@ -1,6 +1,6 @@
 from astropy.cosmology import Planck18 as cosmo  # noqa
 from inspect import isfunction
-from redback.utils import logger, citation_wrapper, calc_ABmag_from_flux_density, lambda_to_nu
+from redback.utils import logger, citation_wrapper, calc_ABmag_from_flux_density, lambda_to_nu, bands_to_frequency
 from redback.constants import day_to_s, speed_of_light, solar_mass, proton_mass, electron_mass, sigma_T
 from redback.sed import get_correct_output_format_from_spectra
 import astropy.units as uu
@@ -285,16 +285,16 @@ class RedbackAfterglows():
         te = np.arcsin(cs / (self.cc * (G2 - 1.) ** 0.5))  # equivalent form for angle due to spreading
         # prepare ex and OmG in this function
         if self.is_expansion:
-            ex = te  # expansion
+            ex = te / (G2) # expansion
             fac = 0.5 * latstep
-            OmG = rotstep * (np.cos(thi - fac) - np.cos(ex + thi + fac))  # equivalent form for linear spacing
+            OmG = rotstep * (np.cos(thi - fac) - np.cos(ex/self.res + thi + fac))  # equivalent form for linear spacing
         else:
-            te = np.ones(te.size)  # no expansion
+            ex = np.ones(te.size)  # no expansion
             fac = 0.5 * latstep
             OmG = rotstep * (np.cos(thi - fac) - np.cos(thi + fac))  # equivalent form for linear spacing
         # prepare R
         size = G.size
-        exponent = ((1 - np.cos(te[0])) / (1 - np.cos(te[:size]))) ** 0.5
+        exponent = ((1 - np.cos(latstep + ex[0])) / (1 - np.cos(latstep + ex[:size]))) ** (1/2)
         R = ((3. - k) * Ne[:size] / (self.fourpi * n)) ** (1. / (3. - k))
         R[1:] = np.diff(R) * exponent[1:size] ** (1. / (3. - k))
         R = np.cumsum(R)
@@ -688,150 +688,10 @@ def _get_kn_dynamics(n0, Eej, Mej):
     # convert from lab frame to observer time. Approximate expression acounting for radial+azimuthal time-of-flight effect
     # (eq. 26 from Nava et al. 2013; see also Waxman 1997)
     tobs = R / (Gamma ** 2 * (1.0 + beta) * speed_of_light) + np.insert(
-        integrate.cumtrapz((1.0 - beta_sh) / beta_sh, x=R) / speed_of_light, 0, 0.0)
+        integrate.cumulative_trapezoid((1.0 - beta_sh) / beta_sh, x=R) / speed_of_light, 0, 0.0)
 
     return t, R, beta, Gamma, eden, tobs, beta_sh, Gamma_sh
 
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2022MNRAS.516.4949S/abstract')
-def kilonova_afterglow_redback(time, redshift, loge0, mej, logn0, logepse, logepsb, p,
-                             **kwargs):
-    """
-    Calculate the afterglow emission from a kilonova remnant, following the model of Sarin et al. 2022.
-    This model was modified by Nikhil Sarin following code provided by Ben Margalit.
-
-    :param time: time in observer frame (days) in observer frame
-    :param redshift: source redshift
-    :param loge0: log10 of the initial kinetic energy of the ejecta (erg)
-    :param mej: ejecta mass (solar masses)
-    :param logn0: log10 of the circumburst density (cm^-3)
-    :param logepse: log10 of the fraction of shock energy given to electrons
-    :param logepsb: log10 of the fraction of shock energy given to magnetic field
-    :param p: power-law index of the electron energy distribution
-    :param kwargs: Additional keyword arguments
-    :param zeta_e: fraction of electrons participating in diffusive shock acceleration. Default is 1.
-    :param output_format: Whether to output flux density or AB mag
-    :param frequency: frequency in Hz for the flux density calculation
-    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
-    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
-        For a proper calculation of the magntitude use the sed variant models.
-    """
-    Eej = 10 ** loge0
-    Mej = mej * solar_mass
-    cosmology = kwargs.get('cosmology', cosmo)
-    epsilon_e = 10 ** logepse
-    epsilon_B = 10 ** logepsb
-    n0 = 10 ** logn0
-    zeta_e = kwargs.get('zeta_e', 1.0)
-    qe = 4.803e-10
-
-    dl = cosmology.luminosity_distance(redshift).cgs.value
-    # calculate blast-wave dynamics
-    t, R, beta, Gamma, eden, tobs, beta_sh, Gamma_sh = _get_kn_dynamics(n0=n0, Eej=Eej, Mej=Mej)
-
-    # shock-amplified magnetic field, minimum & cooling Lorentz factors
-    B = (8.0 * np.pi * epsilon_B * eden) ** 0.5
-    gamma_m = 1.0 + (epsilon_e / zeta_e) * ((p - 2.0) / (p - 1.0)) * (proton_mass / electron_mass) * (Gamma - 1.0)
-    gamma_c = 6.0 * np.pi * electron_mass * speed_of_light / (sigma_T * Gamma * t * B ** 2)
-
-    # number of emitting electrons, where zeta_DN is an approximate smooth interpolation between the "standard"
-    # and deep-Newtonian regime discussed by Sironi & Giannios (2013)
-    zeta_DN = (gamma_m - 1.0) / gamma_m
-    Ne = zeta_DN * zeta_e * (4.0 * np.pi / 3.0) * R ** 3 * n0
-
-    # LOS approximation
-    mu = 1.0
-    blueshift = Gamma * (1.0 - beta * mu)
-
-    frequency = kwargs['frequency']
-    if isinstance(frequency, float):
-        frequency = np.ones(len(time)) * frequency
-    fnu_func = {}
-    for nu in frequency:
-        Fnu_opt_thin = _pnu_synchrotron(nu * blueshift * (1.0 + redshift), B, gamma_m, gamma_c, Ne, p) * (1.0 + redshift) / (
-                    4.0 * np.pi * dl ** 2 * blueshift)
-
-        # correct for synchrotron self-absorption (approximate treatment, correct up to factors of order unity)
-        Fnu_opt_thick = Gamma * (8 * np.pi ** 2 * (nu * blueshift * (1.0 + redshift)) ** 2 / speed_of_light ** 2) * R ** 2 * (
-                    1.0 / 3.0) * electron_mass * speed_of_light ** 2 * np.maximum(gamma_m, (
-                    2 * np.pi * electron_mass * speed_of_light * nu * blueshift * (1.0 + redshift) / (qe * B)) ** 0.5) * (1.0 + redshift) / (
-                                    4.0 * np.pi * dl ** 2 * blueshift)
-        # new prescription:
-        Fnu = Fnu_opt_thick * (1e0 - np.exp(-Fnu_opt_thin / Fnu_opt_thick))
-        # add brute-force optically-thin case to avoid roundoff error in 1e0-np.exp(-x) term (above) when x->0
-        Fnu[Fnu == 0.0] = Fnu_opt_thin[Fnu == 0.0]
-
-        fnu_func[nu] = interp1d(tobs/day_to_s, Fnu, bounds_error=False, fill_value='extrapolate')
-
-    # interpolate onto actual observed frequency and time values
-    flux_density = []
-    for freq, tt in zip(frequency, time):
-        flux_density.append(fnu_func[freq](tt))
-
-    fmjy = np.array(flux_density) / 1e-26
-    if kwargs['output_format'] == 'flux_density':
-        return fmjy
-    elif kwargs['output_format'] == 'magnitude':
-        return calc_ABmag_from_flux_density(fmjy).value
-
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2011Natur.478...82N/abstract')
-def kilonova_afterglow_nakarpiran(time, redshift, loge0, mej, logn0, logepse, logepsb, p, **kwargs):
-    """
-    A kilonova afterglow model based on Nakar & Piran 2011
-
-    :param time: time in days in the observer frame
-    :param redshift: source redshift
-    :param loge0: initial kinetic energy in erg of ejecta
-    :param mej: mass of ejecta in solar masses
-    :param logn0: log10 of the number density of the circumburst medium in cm^-3
-    :param logepse: log10 of the fraction of energy given to electrons
-    :param logepsb: log10 of the fraction of energy given to the magnetic field
-    :param p: electron power law index
-    :param kwargs: Additional keyword arguments
-    :param output_format: Whether to output flux density or AB mag
-    :param frequency: frequency in Hz for the flux density calculation
-    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
-    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
-        For a proper calculation of the magntitude use the sed variant models.
-    :return:
-    """
-    Eej = 10 ** loge0
-    Mej = mej * solar_mass
-    Gamma0 = 1.0 + Eej / (Mej * speed_of_light ** 2)
-    vej = speed_of_light * (1.0 - Gamma0 ** (-2)) ** 0.5
-    cosmology = kwargs.get('cosmology', cosmo)
-    epsilon_e = 10 ** logepse
-    epsilon_B = 10 ** logepsb
-    n0 = 10 ** logn0
-    dl = cosmology.luminosity_distance(redshift).cgs.value
-
-    # in days
-    t_dec = 30 * (Eej / 1e49) ** (1.0 / 3.0) * (n0 / 1e0) ** (-1.0 / 3.0) * (vej / speed_of_light) ** (-5.0 / 3.0)
-
-    fnu_dec_dict = {}
-    fnu_func = {}
-    temp_time = np.linspace(0.1, 100, 200) * t_dec
-    frequency = kwargs['frequency']
-    if isinstance(frequency, float):
-        frequency = np.ones(len(time)) * frequency
-    for freq in frequency:
-        # Eq. 11 in Nakar & Piran 2011 (in Mjy)
-        fnu_dec_dict[freq] = 0.3 * (Eej / 1e49) * n0 ** (0.25 * (p + 1)) * (epsilon_B / 1e-1) ** (0.25 * (p + 1)) * (
-                epsilon_e / 1e-1) ** (p - 1) * (vej / speed_of_light) ** (0.5 * (5 * p - 7)) * (freq / 1.4e9) ** (
-                          -0.5 * (p - 1)) * (dl / 1e27) ** (-2)
-        fnu = fnu_dec_dict[freq] * (temp_time / t_dec) ** 3
-        fnu[temp_time > t_dec] = fnu_dec_dict[freq] * (temp_time[temp_time > t_dec] / t_dec) ** (-0.3 * (5 * p - 7))
-        fnu_func[freq] = interp1d(temp_time, fnu, bounds_error=False, fill_value='extrapolate')
-
-    # interpolate onto actual observed frequency and time values
-    flux_density = []
-    for freq, tt in zip(frequency, time):
-        flux_density.append(fnu_func[freq](tt))
-    fmjy = flux_density * uu.mJy
-    if kwargs['output_format'] == 'flux_density':
-        return fmjy.value
-    elif kwargs['output_format'] == 'magnitude':
-        return calc_ABmag_from_flux_density(fmjy.value).value
-
 @citation_wrapper('redback, https://ui.adsabs.harvard.edu/abs/2018MNRAS.481.2581L/abstract')
 def tophat_redback(time, redshift, thv, loge0, thc, logn0, p, logepse, logepsb, g0, xiN, **kwargs):
     """
@@ -2092,6 +1952,48 @@ def tophat(time, redshift, thv, loge0, thc, logn0, p, logepse, logepsb, ksin, g0
     elif kwargs['output_format'] == 'magnitude':
         return calc_ABmag_from_flux_density(flux_density).value
 
+
+def tophat_from_emulator(time, redshift, thv, loge0, thc, logn0, p, logepse, logepsb, g0, **kwargs):
+    """
+    Evaluate a tophat afterglow model using an mpl regressor. Note that this model predicts for a fixed redshift = 0.01 and fixed ksin = 1.
+    This tophat model does not include all of the ususal kwargs
+
+    :param time: time in days in observer frame, should be in range 0.1 to 300
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param logn0: log10 number density of ISM in cm^-3
+    :param p: electron distribution power law index. Must be greater than 2.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param kwargs: Additional keyword arguments
+    :param frequency: frequency of the band to view in- single number or same length as time array
+    :param output_format: Whether to output flux density or AB mag, specified by 'flux_density' or 'magnitude'
+    :return: flux density or AB mag predicted by emulator. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+        For a proper calculation of the magntitude use the sed variant models
+    """
+
+    from redback_surrogates.afterglowmodels import tophat_emulator
+    
+    z1=0.01
+    z2= redshift
+    frequency= np.log10(kwargs['frequency'])    
+    flux_density = tophat_emulator(new_time=time/(1+z2), thv=thv, loge0=loge0, thc=thc, logn0=logn0, p=p,
+                                            logepse=logepse, logepsb=logepsb, g0=g0,frequency=frequency)
+        
+    #scaling flux density with redshift
+    dl1 = cosmo.luminosity_distance(z1)
+    dl2 = cosmo.luminosity_distance(z2)
+    scale_factor = ((dl1**2)*(1+z1)) / (dl2**2)
+    flux_density=flux_density*scale_factor
+    
+    if kwargs['output_format'] == 'flux_density':
+        return flux_density
+    elif kwargs['output_format'] == 'magnitude':
+        return calc_ABmag_from_flux_density(flux_density).value
+
 @citation_wrapper('https://ui.adsabs.harvard.edu/abs/2020ApJ...896..166R/abstract')
 def afterglow_models_with_energy_injection(time, **kwargs):
     """
@@ -2211,4 +2113,108 @@ def afterglow_models_sed(time, **kwargs):
                                                       spectra=spectra, lambda_array=lambda_observer_frame,
                                                       **kwargs)
 
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_tophat(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, **kwargs):
+    """
+    A tophat jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_tophat(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_tophat(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value
+
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_gaussian(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, **kwargs):
+    """
+    A gaussian jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_gaussian(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_gaussian(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value
 
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2024ApJS..273...17W/abstract')
+def jetsimpy_powerlaw(time, redshift, thv, loge0, thc, nism, A, p, logepse, logepsb, g0, s, **kwargs):
+    """
+    A power-law jet model from jetsimpy
+    :param time: time in days in observer frame
+    :param redshift: source redshift
+    :param thv: viewing angle in radians
+    :param loge0: log10 on axis isotropic equivalent energy
+    :param thc: half width of jet core/jet opening angle in radians
+    :param nism: number density of ISM in cm^-3 (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param A: wind density scale (ntot = A * (r / 1e17)^-2 + nism (cm^-3))
+    :param p: electron distribution power law index.
+    :param logepse: log10 fraction of thermal energy in electrons
+    :param logepsb: log10 fraction of thermal energy in magnetic field
+    :param g0: initial lorentz factor
+    :param s: power-law jet slope
+    :param kwargs: Additional keyword arguments
+    :param output_format: Whether to output flux density or AB mag
+    :param frequency: frequency in Hz for the flux density calculation
+    :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object.
+    :return: flux density or AB mag. Note this is going to give the monochromatic magnitude at the effective frequency for the band.
+    """
+    import jetsimpy #Can not use models unless jetsimpy is downloaded
+    time = time * day_to_s
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    P = dict(Eiso = 10 ** loge0, lf = g0, theta_c = thc, n0 = nism, A = A, eps_e = 10 ** logepse, eps_b = 10 ** logepsb, p = p, theta_v = thv, d = dl*3.24078e-25, z = redshift, s = s) #make a param dict
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        flux_density = jetsimpy.FluxDensity_powerlaw(time, frequency, P)
+        return flux_density   
+    else:
+        frequency = bands_to_frequency(kwargs['bands'])       
+        flux_density = jetsimpy.FluxDensity_powerlaw(time, frequency, P)
+        return calc_ABmag_from_flux_density(flux_density).value

redback/transient_models/combined_models.py

@@ -1,5 +1,11 @@
-import redback.transient_models.extinction_models
-def tophat_and_twolayerstratified(time, redshift, av, thv, loge0, thc, logn0, p, logepse, logepsb, ksin, g0, mej, vej_1, vej_2, kappa, beta, **kwargs):
+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 citation_wrapper
+
+@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,
+                                  logepsb, ksin, g0, mej, vej_1, vej_2, kappa, beta, **kwargs):
     
     """
     function to combine the flux density signals of a tophat afterglow and a two layer stratified kilonova with extinction
@@ -34,20 +40,22 @@ def tophat_and_twolayerstratified(time, redshift, av, thv, loge0, thc, logn0, p,
     :return: flux density signal with extinction added
     
     """
-    kwargs['output_format']='flux_density'
-    kwargs['base_model']='tophat'
-    afterglow = redback.transient_models.extinction_models.extinction_with_afterglow_base_model(time=time, redshift=redshift, av=av,
-         thv=thv, loge0=loge0 , thc= thc, logn0=logn0, p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0,
-        **kwargs)
-    kwargs['base_model']='two_layer_stratified_kilonova'
-    kilonova = redback.transient_models.extinction_models.extinction_with_kilonova_base_model(time=time, redshift=redshift, av=av,
-          mej=mej, vej_1=vej_1, vej_2=vej_2, kappa=kappa, beta=beta, **kwargs)
-    
+    kwargs['output_format'] = 'flux_density'
+    afterglow = tm.afterglow_models.tophat(time=time, redshift=redshift, thv=thv, loge0=loge0, thc=thc, logn0=logn0,
+                                           p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0, **kwargs)
+    kilonova = tm.kilonova_models.two_layer_stratified_kilonova(time=time, redshift=redshift, mej=mej, vej_1=vej_1,
+                                                                vej_2=vej_2, kappa=kappa, beta=beta, **kwargs)
     combined = afterglow+kilonova
+    r_v = kwargs.get('r_v', 3.1)
+    # correct for extinction
+    angstroms = nu_to_lambda(kwargs['frequency'])
+    combined = em._perform_extinction(flux_density=combined, angstroms=angstroms, av=av, r_v=r_v)
     return combined
 
-
-def tophat_and_twocomponent(time, redshift, av, thv, loge0, thc, logn0, p, logepse, logepsb, ksin, g0, mej_1, vej_1, temperature_floor_1, kappa_1, mej_2, vej_2, temperature_floor_2, kappa_2, **kwargs):
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2020ApJ...896..166R/abstract, redback')
+def tophat_and_twocomponent(time, redshift, av, thv, loge0, thc, logn0,
+                            p, logepse, logepsb, ksin, g0, mej_1, vej_1,
+                            temperature_floor_1, kappa_1, mej_2, vej_2, temperature_floor_2, kappa_2, **kwargs):
     
     """
     function to combine the flux density signals of a tophat afterglow and a two component kilonova with extinction added
@@ -85,19 +93,22 @@ def tophat_and_twocomponent(time, redshift, av, thv, loge0, thc, logn0, p, logep
     :return: flux density signal with extinction added
     """
     
-    kwargs['output_format']='flux_density'
-    kwargs['base_model']='tophat'
-    afterglow = redback.transient_models.extinction_models.extinction_with_afterglow_base_model(time=time, redshift=redshift, av=av,
-          thv=thv, loge0=loge0 , thc= thc, logn0=logn0, p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0,
-        **kwargs)
-    kwargs['base_model']='two_component_kilonova_model'
-    kilonova = redback.transient_models.extinction_models.extinction_with_kilonova_base_model(time=time, redshift=redshift, av=av,
-          mej_1=mej_1, vej_1=vej_2, temperature_floor_1=temperature_floor_1, kappa_1=kappa_1,             mej_2=mej_2, vej_2=vej_2, temperature_floor_2=temperature_floor_2, kappa_2=kappa_2, **kwargs)
+    kwargs['output_format'] = 'flux_density'
+    afterglow = tm.afterglow_models.tophat(time=time, redshift=redshift, thv=thv, loge0=loge0, thc=thc, logn0=logn0,
+                                           p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0, **kwargs)
+    kilonova = tm.kilonova_models.two_component_kilonova_model(time=time, redshift=redshift, av=av,
+                                                      mej_1=mej_1, vej_1=vej_1, temperature_floor_1=temperature_floor_1,
+                                                      kappa_1=kappa_1, mej_2=mej_2, vej_2=vej_2,
+                                                      temperature_floor_2=temperature_floor_2, kappa_2=kappa_2, **kwargs)
     
     combined = afterglow + kilonova
+    r_v = kwargs.get('r_v', 3.1)
+    # correct for extinction
+    angstroms = nu_to_lambda(kwargs['frequency'])
+    combined = em._perform_extinction(flux_density=combined, angstroms=angstroms, av=av, r_v=r_v)
     return combined
 
-
+@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2020ApJ...896..166R/abstract, https://ui.adsabs.harvard.edu/abs/1982ApJ...253..785A/abstract')
 def tophat_and_arnett(time, av, redshift, thv, loge0, thc, logn0, p, logepse, logepsb, ksin, g0, f_nickel, mej, **kwargs):
     
     """
@@ -137,46 +148,56 @@ def tophat_and_arnett(time, av, redshift, thv, loge0, thc, logn0, p, logepse, lo
     :return: flux density with extinction added
     """
     
-    kwargs['output_format']='flux_density'
-    kwargs['base_model']='tophat'
-    afterglow = redback.transient_models.extinction_models.extinction_with_afterglow_base_model(time=time, redshift=redshift, av=av,
-         thv=thv, loge0=loge0 , thc= thc, logn0=logn0, p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0,
-        **kwargs)
-    kwargs['base_model']='arnett'
-    supernova = redback.transient_models.extinction_models.extinction_with_supernova_base_model(time=time, redshift=redshift, av=av,
-         f_nickel=f_nickel, mej=mej, **kwargs)
-    
+    kwargs['output_format'] = 'flux_density'
+    afterglow = tm.afterglow_models.tophat(time=time, redshift=redshift, thv=thv, loge0=loge0, thc=thc, logn0=logn0,
+                                           p=p, logepse=logepse, logepsb=logepsb, ksin=ksin, g0=g0, **kwargs)
+    kwargs['base_model'] = 'arnett'
+    supernova = tm.supernova_models.arnett(time=time, redshift=redshift, f_nickel=f_nickel, mej=mej, **kwargs)
     combined = afterglow + supernova
+    r_v = kwargs.get('r_v', 3.1)
+    # correct for extinction
+    angstroms = nu_to_lambda(kwargs['frequency'])
+    combined = em._perform_extinction(flux_density=combined, angstroms=angstroms, av=av, r_v=r_v)
     return combined
 
-def afterglow_and_optical(time, redshift, av, model_type, afterglow_kwargs, optical_kwargs, **shared_kwargs):
-    
+@citation_wrapper('redback, and any citations for the specific model you use')
+def afterglow_and_optical(time, redshift, av, **model_kwargs):
     """
     function to combine the signals of any afterglow and any other optical transient with extinction added
     
     :param time: time in days in observer frame
     :param redshift: source redshift
     :param av: absolute mag extinction
-    :param model_type: specify type of optical transient model- 'supernova', 'tde', 'kilonova', 'magnetar_driven', 'shock_powered'
+    :param model_kwargs: kwargs shared by models frequency and r_v (extinction paramater 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 optical_kwargs: dictionary of parameters required by the optical transient model specifed 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 shared_kwargs: kwargs shared by models e.g. output_format, frequency, bands, r_v (extinction paramater defaults to 3.1)
-    :return: set by shared_kwargs output format - 'flux_density' or 'magnitude' with extinction added
-        note that only afterglow_models_sed allow for magnitude outputs
+    :return: flux density signal with extinction added
     """
-    
-    afterglow_kwargs.update(shared_kwargs)
-    optical_kwargs.update(shared_kwargs)
-   
-    afterglow = redback.transient_models.extinction_models._evaluate_extinction_model(time=time, redshift=redshift, av=av,
-                                                                                   model_type='afterglow', **afterglow_kwargs)
-    
-    optical= redback.transient_models.extinction_models._evaluate_extinction_model(time=time, redshift=redshift, av=av,
-                                                                   model_type=model_type, **optical_kwargs)
-  
-    combined= afterglow + optical
+
+    from redback.model_library import all_models_dict
+    optical_kwargs = model_kwargs['optical_kwargs']
+    afterglow_kwargs = model_kwargs['afterglow_kwargs']
+    model_kwargs['output_format']= model_kwargs.get('output_format', 'flux_density')
+
+    _afterglow_kwargs = afterglow_kwargs.copy()
+    _afterglow_kwargs.update(model_kwargs)
+
+    _optical_kwargs = optical_kwargs.copy()
+    _optical_kwargs.update(model_kwargs)
+
+    afterglow_function = all_models_dict[_afterglow_kwargs['base_model']]
+    afterglow = afterglow_function(time=time, redshift=redshift,  **_afterglow_kwargs)
+
+    optical_function = all_models_dict[_optical_kwargs['base_model']]
+    optical = optical_function(time=time, redshift=redshift, **_optical_kwargs)
+
+    combined = afterglow + optical
+    r_v = model_kwargs.get('r_v', 3.1)
+    # correct for extinction
+    angstroms = nu_to_lambda(model_kwargs['frequency'])
+    combined = em._perform_extinction(flux_density=combined, angstroms=angstroms, av=av, r_v=r_v)
     return combined
     

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,30 +11,32 @@ extinction_afterglow_base_models = ['tophat', 'cocoon', 'gaussian',
                                     'kn_afterglow', 'cone_afterglow',
                                     'gaussiancore', 'gaussian',
                                     'smoothpowerlaw', 'powerlawcore',
+                                    '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',
                                     'csm_interaction', 'csm_nickel', 'type_1a', 'type_1c',
-                                    'general_magnetar_slsn','general_magnetar_driven_supernova']
+                                    'general_magnetar_slsn','general_magnetar_driven_supernova', 'sn_fallback']
 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_cooling_envelope',
-                              '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']
 
@@ -86,7 +87,6 @@ def _perform_extinction(flux_density, angstroms, av, r_v):
     :return: flux
     """
     import extinction  # noqa
-    import numpy.ma as ma
     if isinstance(angstroms, float):
         angstroms = np.array([angstroms])    
     mag_extinction = extinction.fitzpatrick99(angstroms, av, r_v=r_v)