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

redback/sed.py

@@ -6,6 +6,193 @@ from sncosmo import TimeSeriesSource
 from redback.constants import *
 from redback.utils import nu_to_lambda, bandpass_magnitude_to_flux
 
+def _bandflux_single_redback(model, band, time_or_phase):
+    """
+
+    Synthetic photometry of a model through a single bandpass
+
+    :param model: Source object
+    :param band: Bandpass
+    :param time_or_phase: Time or phase numpy array
+    :return: bandflux through the bandpass
+    """
+    from sncosmo.utils import integration_grid
+    HC_ERG_AA = 1.9864458571489284e-08 # planck * speed_of_light in AA/s
+    MODEL_BANDFLUX_SPACING = 5.0  # Angstroms
+
+    if (band.minwave() < model.minwave() or band.maxwave() > model.maxwave()):
+        raise ValueError('bandpass {0!r:s} [{1:.6g}, .., {2:.6g}] '
+                         'outside spectral range [{3:.6g}, .., {4:.6g}]'
+                         .format(band.name, band.minwave(), band.maxwave(),
+                                 model.minwave(), model.maxwave()))
+
+        # Set up wavelength grid. Spacing (dwave) evenly divides the bandpass,
+        # closest to 5 angstroms without going over.
+    wave, dwave = integration_grid(band.minwave(), band.maxwave(),
+                                   MODEL_BANDFLUX_SPACING)
+    trans = band(wave)
+    f = model._flux(time_or_phase, wave)
+    f = np.abs(f)
+    return np.sum(wave * trans * f, axis=1) * dwave / HC_ERG_AA
+
+
+def _bandflux_redback(model, band, time_or_phase, zp, zpsys):
+    """
+    Support function for bandflux in Source and Model. Follows SNCOSMO
+
+    This is necessary to have outside because ``phase`` is used in Source
+    and ``time`` is used in Model, and we want the method signatures to
+    have the right variable name.
+    """
+    from sncosmo.magsystems import get_magsystem
+    from sncosmo.bandpasses import get_bandpass
+
+    if zp is not None and zpsys is None:
+        raise ValueError('zpsys must be given if zp is not None')
+
+    # broadcast arrays
+    if zp is None:
+        time_or_phase, band = np.broadcast_arrays(time_or_phase, band)
+    else:
+        time_or_phase, band, zp, zpsys = \
+            np.broadcast_arrays(time_or_phase, band, zp, zpsys)
+
+    # Convert all to 1-d arrays.
+    ndim = time_or_phase.ndim  # Save input ndim for return val.
+    time_or_phase = np.atleast_1d(time_or_phase)
+    band = np.atleast_1d(band)
+    if zp is not None:
+        zp = np.atleast_1d(zp)
+        zpsys = np.atleast_1d(zpsys)
+
+    # initialize output arrays
+    bandflux = np.zeros(time_or_phase.shape, dtype=float)
+
+    # Loop over unique bands.
+    for b in set(band):
+        mask = band == b
+        b = get_bandpass(b)
+
+        fsum = _bandflux_single_redback(model, b, time_or_phase[mask])
+
+        if zp is not None:
+            zpnorm = 10. ** (0.4 * zp[mask])
+            bandzpsys = zpsys[mask]
+            for ms in set(bandzpsys):
+                mask2 = bandzpsys == ms
+                ms = get_magsystem(ms)
+                zpnorm[mask2] = zpnorm[mask2] / ms.zpbandflux(b)
+            fsum *= zpnorm
+
+        bandflux[mask] = fsum
+
+    if ndim == 0:
+        return bandflux[0]
+    return bandflux
+
+def _bandmag_redback(model, band, magsys, time_or_phase):
+    """
+    Support function for bandflux in Source and Model.
+    This is necessary to have outside the models because ``phase`` is used in
+    Source and ``time`` is used in Model.
+    """
+    from sncosmo.magsystems import get_magsystem
+
+    bandflux = _bandflux_redback(model, band, time_or_phase, None, None)
+    band, magsys, bandflux = np.broadcast_arrays(band, magsys, bandflux)
+    return_scalar = (band.ndim == 0)
+    band = band.ravel()
+    magsys = magsys.ravel()
+    bandflux = bandflux.ravel()
+
+    result = np.empty(bandflux.shape, dtype=float)
+    for i, (b, ms, f) in enumerate(zip(band, magsys, bandflux)):
+        ms = get_magsystem(ms)
+        zpf = ms.zpbandflux(b)
+        result[i] = -2.5 * np.log10(f / zpf)
+
+    if return_scalar:
+        return result[0]
+    return result
+
+class RedbackTimeSeriesSource(TimeSeriesSource):
+        def __init__(self, phase, wave, flux, **kwargs):
+            """
+            RedbackTimeSeriesSource is a subclass of sncosmo.TimeSeriesSource that adds the ability to return the
+            flux density at a given time and wavelength, and changes
+            the behaviour of the _flux method to better handle models with very low flux values.
+
+            :param phase: phase/time array
+            :param wave: wavelength array in Angstrom
+            :param spectra: spectra in erg/cm^2/s/A evaluated at all times and frequencies; shape (len(times), len(frequency_array))
+            :param kwargs: additional arguments
+            """
+            super(RedbackTimeSeriesSource, self).__init__(phase=phase, wave=wave, flux=flux, **kwargs)
+
+        def get_flux_density(self, time, wavelength):
+            """
+            Get the flux density at a given time and wavelength.
+
+            :param time: time in days
+            :param wavelength: wavelength in Angstrom
+            :return: flux density in erg/cm^2/s/A
+            """
+            return self._flux(time, wavelength)
+
+        def bandflux(self, band, phase, zp=None, zpsys=None):
+            """
+            Flux through the given bandpass(es) at the given phase(s).
+
+            Default return value is flux in photons / s / cm^2. If zp and zpsys
+            are given, flux(es) are scaled to the requested zeropoints.
+
+            Parameters
+            ----------
+            band : str or list_like
+                Name(s) of bandpass(es) in registry.
+            phase : float or list_like, optional
+                Phase(s) in days. Default is `None`, which corresponds to the full
+                native phase sampling of the model.
+            zp : float or list_like, optional
+                If given, zeropoint to scale flux to (must also supply ``zpsys``).
+                If not given, flux is not scaled.
+            zpsys : str or list_like, optional
+                Name of a magnitude system in the registry, specifying the system
+                that ``zp`` is in.
+
+            Returns
+            -------
+            bandflux : float or `~numpy.ndarray`
+                Flux in photons / s /cm^2, unless `zp` and `zpsys` are
+                given, in which case flux is scaled so that it corresponds
+                to the requested zeropoint. Return value is `float` if all
+                input parameters are scalars, `~numpy.ndarray` otherwise.
+            """
+            return _bandflux_redback(self, band, phase, zp, zpsys)
+
+        def bandmag(self, band, magsys, phase):
+            """Magnitude at the given phase(s) through the given
+            bandpass(es), and for the given magnitude system(s).
+
+            Parameters
+            ----------
+            band : str or list_like
+                Name(s) of bandpass in registry.
+            magsys : str or list_like
+                Name(s) of `~sncosmo.MagSystem` in registry.
+            phase : float or list_like
+                Phase(s) in days.
+
+            Returns
+            -------
+            mag : float or `~numpy.ndarray`
+                Magnitude for each item in band, magsys, phase.
+                The return value is a float if all parameters are not iterables.
+                The return value is an `~numpy.ndarray` if any are iterable.
+            """
+            return _bandmag_redback(self, band, magsys, phase)
+
+
 
 def blackbody_to_flux_density(temperature, r_photosphere, dl, frequency):
     """
@@ -298,13 +485,18 @@ def get_correct_output_format_from_spectra(time, time_eval, spectra, lambda_arra
     :param time: times in observer frame in days to evaluate the model on
     :param time_eval: times in observer frame where spectra are evaluated. A densely sampled array for accuracy
     :param bands: band array - must be same length as time array or a single band
-    :param spectra: spectra in mJy evaluated at all times and frequencies; shape (len(times), len(frequency_array))
+    :param spectra: spectra in erg/cm^2/s/A evaluated at all times and frequencies; shape (len(times), len(frequency_array))
     :param lambda_array: wavelenth array in Angstrom in observer frame
     :param kwargs: Additional parameters
     :param output_format: 'flux', 'magnitude', 'sncosmo_source', 'flux_density'
     :return: flux, magnitude or SNcosmo TimeSeries Source depending on output format kwarg
     """
-    source = TimeSeriesSource(phase=time_eval, wave=lambda_array, flux=spectra)
+    # clean up spectrum to remove nonsensical values before creating sncosmo source
+    spectra = np.nan_to_num(spectra)
+    spectra[spectra.value == np.nan_to_num(np.inf)] = 1e-30 * np.mean(spectra[5])
+    spectra[spectra.value == 0.] = 1e-30 * np.mean(spectra[5])
+
+    source = RedbackTimeSeriesSource(phase=time_eval, wave=lambda_array, flux=spectra)
     if kwargs['output_format'] == 'flux':
         bands = kwargs['bands']
         magnitude = source.bandmag(phase=time, band=bands, magsys='ab')

redback/simulate_transients.py

@@ -14,7 +14,7 @@ datadir = os.path.join(os.path.dirname(redback.__file__), 'tables')
 
 class SimulateGenericTransient(object):
     def __init__(self, model, parameters, times, model_kwargs, data_points,
-                 seed=1234, multiwavelength_transient=False, noise_term=0.2):
+                 seed=1234, multiwavelength_transient=False, noise_term=0.2, noise_type='gaussianmodel', extra_scatter=0.0):
         """
         A generic interface to simulating transients
 
@@ -31,7 +31,12 @@ class SimulateGenericTransient(object):
                 and the data points are sampled in bands/frequency as well,
                 rather than just corresponding to one wavelength/filter.
                 This also allows the same time value to be sampled multiple times.
-        :param noise_term: Float. Factor which is multiplied by the model flux/magnitude to give the sigma.
+        :param noise_type: String. Type of noise to add to the model.
+            Default is 'gaussianmodel' where sigma is noise_term * model.
+            Another option is 'gaussian' i.e., a simple Gaussian noise with sigma = noise_term.
+        :param noise_term: Float. Factor which is multiplied by the model flux/magnitude to give the sigma
+            or is sigma itself for 'gaussian' noise.
+        :param extra_scatter: Float. Sigma of normal added to output for additional scatter.
         """
         self.model = redback.model_library.all_models_dict[model]
         self.parameters = parameters
@@ -44,20 +49,25 @@ class SimulateGenericTransient(object):
         self.noise_term = noise_term
         random.seed(self.seed)
 
-        if multiwavelength_transient:
-            self.all_bands = self.model_kwargs.get('bands', None)
-            self.all_frequency = self.model_kwargs.get('frequency', None)
-            if self.all_bands is None and self.all_frequency is None:
-                raise ValueError('Must supply either bands or frequency to sample data points for an optical transient')
-            if self.all_bands is not None and self.all_frequency is None:
-                self.subset_bands = np.array(random.choices(self.all_bands, k=self.data_points))
-            if self.all_bands is None and self.all_frequency is not None:
-                self.subset_frequency = np.array(random.choices(self.all_frequency, k=self.data_points))
-            self.replacement = True
-            # allow times to be chosen repeatedly
+        self.all_bands = self.model_kwargs.get('bands', None)
+        self.all_frequency = self.model_kwargs.get('frequency', None)
+        if self.all_bands is None and self.all_frequency is None:
+            raise ValueError('Must supply either bands or frequency to sample data points for an optical transient')
         else:
-            # allow times to be chosen only once.
-            self.replacement = False
+            if multiwavelength_transient:
+                if self.all_bands is not None and self.all_frequency is None:
+                    self.subset_bands = np.array(random.choices(self.all_bands, k=self.data_points))
+                if self.all_bands is None and self.all_frequency is not None:
+                    self.subset_frequency = np.array(random.choices(self.all_frequency, k=self.data_points))
+                self.replacement = True
+                # allow times to be chosen repeatedly
+            else:
+                if self.all_bands is not None and self.all_frequency is None:
+                    self.subset_bands = self.data_points * [self.all_bands]
+                if self.all_bands is None and self.all_frequency is not None:
+                    self.subset_frequency = np.ones(self.data_points) * self.all_frequency
+                # allow times to be chosen only once.
+                self.replacement = False
         self.subset_times = np.sort(np.random.choice(self.all_times, size=self.data_points, replace=self.replacement))
 
         injection_kwargs = self.parameters.copy()
@@ -76,24 +86,43 @@ class SimulateGenericTransient(object):
         if 'frequency' in model_kwargs.keys():
             data['frequency'] = self.subset_frequency
         data['true_output'] = true_output
-        data['output'] = np.random.normal(true_output, self.noise_term * true_output)
-        data['output_error'] = self.noise_term * true_output
+
+        if noise_type == 'gaussianmodel':
+            noise = np.random.normal(0, self.noise_term * true_output, len(true_output))
+            output = true_output + noise
+            output_error = self.noise_term * true_output
+        elif noise_type == 'gaussian':
+            noise = np.random.normal(0, self.noise_term, len(true_output))
+            output = true_output + noise
+            output_error = self.noise_term
+        else:
+            logger.warning(f"noise_type {noise_type} not implemented.")
+            raise ValueError('noise_type must be either gaussianmodel or gaussian')
+
+        if extra_scatter > 0:
+            extra_noise = np.random.normal(0, extra_scatter, len(true_output))
+            output = output + extra_noise
+            output_error = np.sqrt(output_error**2 + extra_noise**2)
+
+        data['output'] = output
+        data['output_error'] = output_error
         self.data = data
 
-        def save_transient(self, name):
-            """
-            Save the transient observations to a csv file.
-            This will save the full observational dataframe including non-detections etc.
-            This will save the data to a folder called 'simulated'
-            with the name of the transient and a csv file of the injection parameters
-
-            :param name: name to save transient.
-            """
-            bilby.utils.check_directory_exists_and_if_not_mkdir('simulated')
-            path = 'simulated/' + name + '.csv'
-            injection_path = 'simulated/' + name + '_injection_parameters.csv'
-            self.observations.to_csv(path, index=False)
-            self.parameters.to_csv(injection_path, index=False)
+    def save_transient(self, name):
+        """
+        Save the transient observations to a csv file.
+        This will save the full observational dataframe including non-detections etc.
+        This will save the data to a folder called 'simulated'
+        with the name of the transient and a csv file of the injection parameters
+
+        :param name: name to save transient.
+        """
+        bilby.utils.check_directory_exists_and_if_not_mkdir('simulated')
+        path = 'simulated/' + name + '.csv'
+        injection_path = 'simulated/' + name + '_injection_parameters.csv'
+        self.data.to_csv(path, index=False)
+        self.parameters=pd.DataFrame.from_dict([self.parameters])
+        self.parameters.to_csv(injection_path, index=False)
 
 class SimulateOpticalTransient(object):
     def __init__(self, model, parameters, pointings_database=None,
@@ -415,7 +444,7 @@ class SimulateOpticalTransient(object):
         Convert the circular field of view to a radius in radians.
         :return: survey_radius in radians
         """
-        survey_fov_sqrad = self.survey_fov_sqdeg*(np.pi/180.0)
+        survey_fov_sqrad = self.survey_fov_sqdeg*(np.pi/180.0)**2
         survey_radius = np.sqrt(survey_fov_sqrad/np.pi)
         # survey_radius = np.sqrt(self.survey_fov_sqdeg*((np.pi/180.0)**2.0)/np.pi)
         return survey_radius

redback/tables/filters.csv

@@ -212,7 +212,7 @@ f2100w,1.443e+13,207950.05,black,1.122e-07,f2100w,MIRI/F2100W
 f2300c,1.325e+13,226446.44,black,9.225e-08,f2300c,MIRI/F2300C
 f2550w,1.183e+13,253640.02,black,6.601e-08,f2550w,MIRI/F2550W
 G,4.799e+14,6251.50,black,1.321e-05,gaia::G,Gaia/G
-gaia::G,4.799e+14,6251.50,black,1.321e-05,gaia::G,Gaia/G
+gaia::g,4.799e+14,6251.50,black,1.321e-05,gaia::g,Gaia/G
 y',3.116e+14,9627.79,black,7.517e-07,ps1::y,PS1/y
 y,3.116e+14,9627.79,black,7.517e-07,ps1::y,PS1/y
 Y,3.116e+14,9627.79,black,7.517e-07,ps1::y,PS1/y
@@ -238,3 +238,17 @@ euclid::VIS,4.223e+14,7103.37,black,7.577e-06,euclid::VIS,EUCLID/VIS
 efosc2::g,5.866e+14,5114.66,black,3.302e-06,efosc2::g,EFOSC2/g
 efosc2::r,4.493e+14,6677.40,black,2.342e-06,efosc2::r,EFOSC2/r
 efosc2::i,3.758e+14,7983.20,black,2.435e-06,efosc2::i,EFOSC2/i
+desg,6.2390e+14,4808.49,black,4.781e-06,desg,DES/g
+desr,4.6746e+14,6417.65,black,3.409e-06,desr,DES/r
+desi,3.8390e+14,7814.58,black,2.300e-06,desi,DES/i
+desz,3.2719e+14,9168.85,black,1.679e-06,desz,DES/z
+desy,3.0315e+14,9896.11,black,7.511e-07,desy,DES/y
+tess,3.8973e+14,7697.6,black,7.658e-06,tess,TESS
+gaia::gbp,5.8712e+14,5109.71,black,9.421e-06,gaia::gbp,GAIA/Gbp
+gaia::g,4.799e+14,6217.59,black,1.278e-05,gaia::g,GAIA/G
+gaia::grp,3.8615e+14,7769.02,black,5.472e-06,gaia::grp,GAIA/Grp
+gaia::grvs,3.4970e+14,8578.76,black,3.481e-07,gaia::grvs,GAIA/Grvs
+gotob,6.4895e+14,4622.88,black,3.771e-06,gotob,GOTO/B
+gotog,5.5819e+14,5374.54,black,3.255e-06,gotog,GOTO/G
+gotol,5.4311e+14,5523.78,black,8.027e-06,gotol,GOTO/L
+gotor,4.6679e+14,6426.84,black,2.510e-06,gotor,GOTO/R

redback/transient/afterglow.py

@@ -57,7 +57,7 @@ class Afterglow(Transient):
         :type flux: np.ndarray, optional
         :type flux_err: np.ndarray, optional
         :param flux_err: Flux error values.
-        :param flux_density:Flux density values.
+        :param flux_density: Flux density values.
         :type flux_density: np.ndarray, optional
         :param flux_density_err: Flux density error values.
         :type flux_density_err: np.ndarray, optional
@@ -244,7 +244,8 @@ class Afterglow(Transient):
                 'BAT Photon Index (15-150 keV) (PL = simple power-law, CPL = cutoff power-law)'].fillna(0)
             self.meta_data = meta_data
         except FileNotFoundError:
-            logger.warning("Meta data does not exist for this event.")
+            logger.info("Metadata does not exist for this event.")
+            logger.info("Setting metadata to None. This is not an error, but a warning that no metadata could be found online.")
             self.meta_data = None
 
     def _set_photon_index(self) -> None:

redback/transient/kilonova.py

@@ -50,7 +50,7 @@ class Kilonova(OpticalTransient):
         :type flux: np.ndarray, optional
         :type flux_err: np.ndarray, optional
         :param flux_err: Flux error values.
-        :param flux_density:Flux density values.
+        :param flux_density: Flux density values.
         :type flux_density: np.ndarray, optional
         :param flux_density_err: Flux density error values.
         :type flux_density_err: np.ndarray, optional

redback/transient/supernova.py

@@ -45,7 +45,7 @@ class Supernova(OpticalTransient):
         :type flux: np.ndarray, optional
         :type flux_err: np.ndarray, optional
         :param flux_err: Flux error values.
-        :param flux_density:Flux density values.
+        :param flux_density: Flux density values.
         :type flux_density: np.ndarray, optional
         :param flux_density_err: Flux density error values.
         :type flux_density_err: np.ndarray, optional

redback/transient/tde.py

@@ -42,7 +42,7 @@ class TDE(OpticalTransient):
         :type flux: np.ndarray, optional
         :type flux_err: np.ndarray, optional
         :param flux_err: Flux error values.
-        :param flux_density:Flux density values.
+        :param flux_density: Flux density values.
         :type flux_density: np.ndarray, optional
         :param flux_density_err: Flux density error values.
         :type flux_density_err: np.ndarray, optional

redback/transient/transient.py

@@ -137,7 +137,7 @@ class Transient(object):
         self.system = system
         self.data_mode = data_mode
         self.active_bands = active_bands
-        self.sncosmo_bands = redback.utils.sncosmo_bandname_from_band(self.bands, warning_style='soft')
+        self.sncosmo_bands = redback.utils.sncosmo_bandname_from_band(self.bands)
         self.redshift = redshift
         self.name = name
         self.use_phase_model = use_phase_model
@@ -906,7 +906,7 @@ class OpticalTransient(Transient):
             meta_data = pd.read_csv(self.event_table, on_bad_lines='skip', delimiter=',', dtype='str')
         except FileNotFoundError as e:
             redback.utils.logger.warning(e)
-            redback.utils.logger.warning("Setting metadata to None")
+            redback.utils.logger.warning("Setting metadata to None. This is not an error, but a warning that no metadata could be found online.")
             meta_data = None
         self.meta_data = meta_data
 

redback/transient_models/afterglow_models.py

@@ -2092,6 +2092,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):
     """

redback/transient_models/combined_models.py

@@ -1,5 +1,9 @@
-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
+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
@@ -36,18 +40,20 @@ def tophat_and_twolayerstratified(time, redshift, av, thv, loge0, thc, logn0, p,
     """
     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,
+    afterglow = em.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,
+    kilonova = em.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)
     
     combined = afterglow+kilonova
     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 +91,21 @@ 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'
+    kwargs['base_model'] = 'tophat'
+    afterglow = em.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 = em.extinction_with_kilonova_base_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
     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,19 +145,20 @@ 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['output_format'] = 'flux_density'
+    kwargs['base_model'] = 'tophat'
+    afterglow = em.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,
+    kwargs['base_model'] = 'arnett'
+    supernova = em.extinction_with_supernova_base_model(time=time, redshift=redshift, av=av,
          f_nickel=f_nickel, mej=mej, **kwargs)
     
     combined = afterglow + supernova
     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
@@ -157,26 +166,32 @@ def afterglow_and_optical(time, redshift, av, model_type, afterglow_kwargs, opti
     :param time: time in days in observer frame
     :param redshift: source redshift
     :param av: absolute mag extinction
+    :param model_kwargs: kwargs shared by models e.g. output_format, frequency, bands, r_v (extinction paramater defaults to 3.1)
     :param model_type: specify type of optical transient model- 'supernova', 'tde', 'kilonova', 'magnetar_driven', 'shock_powered'
     :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
     """
+
+    optical_kwargs = model_kwargs['optical_kwargs']
+    afterglow_kwargs = model_kwargs['afterglow_kwargs']
+
+    _afterglow_kwargs = afterglow_kwargs.copy()
+    _afterglow_kwargs.update(model_kwargs)
+    _afterglow_kwargs.pop('model_type')
+
+    _optical_kwargs = optical_kwargs.copy()
+    _optical_kwargs.update(model_kwargs)
+
+    afterglow = em._evaluate_extinction_model(time=time, redshift=redshift, av=av,model_type='afterglow',
+                                              **_afterglow_kwargs)
     
-    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)
-  
+    optical = em._evaluate_extinction_model(time=time, redshift=redshift, av=av, **_optical_kwargs)
+
     combined= afterglow + optical
     return combined