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

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)

redback/transient_models/tde_models.py

@@ -262,6 +262,8 @@ def gaussianrise_cooling_envelope_bolometric(time, peak_time, sigma_t, mbh_6, st
     bolometric version for fitting the bolometric lightcurve
 
     :param time: time in source frame in days
+    :param peak_time: peak time in days
+    :param sigma_t: the sharpness of the Gaussian in days
     :param mbh_6: mass of supermassive black hole in units of 10^6 solar mass
     :param stellar_mass: stellar mass in units of solar masses
     :param eta: SMBH feedback efficiency (typical range: etamin - 0.1)
@@ -299,6 +301,8 @@ def gaussianrise_cooling_envelope(time, redshift, peak_time, sigma_t, mbh_6, ste
 
     :param time: time in observer frame in days
     :param redshift: redshift
+    :param peak_time: peak time in days
+    :param sigma_t: the sharpness of the Gaussian in days
     :param mbh_6: mass of supermassive black hole in units of 10^6 solar mass
     :param stellar_mass: stellar mass in units of solar masses
     :param eta: SMBH feedback efficiency (typical range: etamin - 0.1)
@@ -408,6 +412,127 @@ def gaussianrise_cooling_envelope(time, redshift, peak_time, sigma_t, mbh_6, ste
             output.append(flux_den_interp_func[freq](tt * cc.day_to_s))
         return np.array(output)
 
+@citation_wrapper('https://arxiv.org/abs/2307.15121,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
+def bpl_cooling_envelope(time, redshift, peak_time, alpha_1, alpha_2, mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
+    """
+    Full lightcurve, with gaussian rise till fallback time and then the metzger tde model,
+    photometric version where each band is fit/joint separately
+
+    :param time: time in observer frame in days
+    :param redshift: redshift
+    :param peak_time: peak time in days
+    :param alpha_1: power law index for first power law
+    :param alpha_2: power law index for second power law (should be positive)
+    :param mbh_6: mass of supermassive black hole in units of 10^6 solar mass
+    :param stellar_mass: stellar mass in units of solar masses
+    :param eta: SMBH feedback efficiency (typical range: etamin - 0.1)
+    :param alpha: disk viscosity
+    :param beta: TDE penetration factor (typical range: 1 - beta_max)
+    :param kwargs: Additional parameters
+    :param xi: Optional argument (default set to one) to change the point where lightcurve switches from Gaussian rise to cooling envelope.
+        stitching_point = xi * tfb (where tfb is fallback time). So a xi=1 means the stitching point is at fallback time.
+    :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', 'flux'
+    :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', 'flux'
+    """
+    binding_energy_const = kwargs.get('binding_energy_const', 0.8)
+    tfb_sf = calc_tfb(binding_energy_const, mbh_6, stellar_mass)  # source frame
+    tfb_obf = tfb_sf * (1. + redshift)  # observer frame
+    xi = kwargs.get('xi', 1.)
+    output = _cooling_envelope(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
+    cosmology = kwargs.get('cosmology', cosmo)
+    dl = cosmology.luminosity_distance(redshift).cgs.value
+    stitching_point = xi * tfb_obf
+
+    # normalisation term in observer frame
+    f1 = pm.exponential_powerlaw(time=stitching_point, a_1=1., tpeak=peak_time * cc.day_to_s,
+                                 alpha_1=alpha_1, alpha_2=alpha_2)
+
+    if kwargs['output_format'] == 'flux_density':
+        frequency = kwargs['frequency']
+        if isinstance(frequency, float):
+            frequency = np.ones(len(time)) * frequency
+
+        # convert to source frame time and frequency
+        frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time)
+        unique_frequency = np.sort(np.unique(frequency))
+
+        # source frame
+        f2 = sed.blackbody_to_flux_density(temperature=output.photosphere_temperature[0],
+                                           r_photosphere=output.photosphere_radius[0],
+                                           dl=dl, frequency=unique_frequency).to(uu.mJy)
+        norms = f2.value / f1
+        norm_dict = dict(zip(unique_frequency, norms))
+
+        # build flux density function for each frequency
+        flux_den_interp_func = {}
+        for freq in unique_frequency:
+            tt_pre_fb = np.linspace(0, stitching_point / cc.day_to_s, 200) * cc.day_to_s
+            tt_post_fb = xi * (output.time_temp * (1 + redshift))
+            total_time = np.concatenate([tt_pre_fb, tt_post_fb])
+            f1 = pm.exponential_powerlaw(time=tt_pre_fb, a_1=norm_dict[freq],
+                                         tpeak=peak_time * cc.day_to_s, alpha_1=alpha_1, alpha_2=alpha_2)
+            f2 = sed.blackbody_to_flux_density(temperature=output.photosphere_temperature,
+                                               r_photosphere=output.photosphere_radius,
+                                               dl=dl, frequency=freq).to(uu.mJy)
+            flux_den = np.concatenate([f1, f2.value])
+            flux_den_interp_func[freq] = interp1d(total_time, flux_den, fill_value='extrapolate')
+
+        # interpolate onto actual observed frequency and time values
+        flux_density = []
+        for freq, tt in zip(frequency, time):
+            flux_density.append(flux_den_interp_func[freq](tt * cc.day_to_s))
+        flux_density = flux_density * uu.mJy
+        return flux_density.to(uu.mJy).value
+    else:
+        bands = kwargs['bands']
+        if isinstance(bands, str):
+            bands = [str(bands) for x in range(len(time))]
+
+        unique_bands = np.unique(bands)
+        temp_kwargs = kwargs.copy()
+        temp_kwargs['bands'] = unique_bands
+        f2 = cooling_envelope(time=0., redshift=redshift,
+                              mbh_6=mbh_6, stellar_mass=stellar_mass, eta=eta, alpha=alpha, beta=beta,
+                              **temp_kwargs)
+        if kwargs['output_format'] == 'magnitude':
+            # make the normalisation in fmjy to avoid magnitude normalisation problems
+            _f2mjy = calc_flux_density_from_ABmag(f2).value
+            norms = _f2mjy / f1
+        else:
+            norms = f2 / f1
+
+        if isinstance(norms, float):
+            norms = np.ones(len(time)) * norms
+        norm_dict = dict(zip(unique_bands, norms))
+
+        flux_den_interp_func = {}
+        for band in unique_bands:
+            tt_pre_fb = np.linspace(0, stitching_point / cc.day_to_s, 100) * cc.day_to_s
+            tt_post_fb = output.time_temp * (1 + redshift)
+            total_time = np.concatenate([tt_pre_fb, tt_post_fb])
+            f1 = pm.exponential_powerlaw(time=tt_pre_fb, a_1=norm_dict[band],
+                                         tpeak=peak_time * cc.day_to_s, alpha_1=alpha_1, alpha_2=alpha_2)
+            if kwargs['output_format'] == 'magnitude':
+                f1 = calc_ABmag_from_flux_density(f1).value
+            temp_kwargs = kwargs.copy()
+            temp_kwargs['bands'] = band
+            f2 = cooling_envelope(time=output.time_since_fb / cc.day_to_s, redshift=redshift,
+                                  mbh_6=mbh_6, stellar_mass=stellar_mass, eta=eta, alpha=alpha, beta=beta,
+                                  **temp_kwargs)
+            flux_den = np.concatenate([f1, f2])
+            flux_den_interp_func[band] = interp1d(total_time, flux_den, fill_value='extrapolate')
+
+        # interpolate onto actual observed band and time values
+        output = []
+        for freq, tt in zip(bands, time):
+            output.append(flux_den_interp_func[freq](tt * cc.day_to_s))
+        return np.array(output)
+
 @citation_wrapper('redback')
 def tde_analytical_bolometric(time, l0, t_0_turn, **kwargs):
     """
@@ -471,7 +596,7 @@ def tde_analytical(time, redshift, l0, t_0_turn, **kwargs):
     else:
         time_obs = time
         lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 100))
-        time_temp = np.geomspace(0.1, 300, 200) # in days
+        time_temp = np.geomspace(0.1, 1000, 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)

redback/utils.py

@@ -44,7 +44,7 @@ def download_pointing_tables():
     """
     return logger.info("Pointing tables downloaded and stored in redback/tables")
 
-def sncosmo_bandname_from_band(bands, warning_style='soft'):
+def sncosmo_bandname_from_band(bands, warning_style='softest'):
     """
     Convert redback data band names to sncosmo compatible band names
 
@@ -65,12 +65,14 @@ def sncosmo_bandname_from_band(bands, warning_style='soft'):
         try:
             res.append(bands_to_flux[band])
         except KeyError as e:
-            logger.info(e)
             if warning_style == 'hard':
                 raise KeyError(f"Band {band} is not defined in filters.csv!")
-            else:
+            elif warning_style == 'soft':
+                logger.info(e)
                 logger.info(f"Band {band} is not defined in filters.csv!")
                 res.append('r')
+            else:
+                res.append('r')
     return np.array(res)
 
 def check_kwargs_validity(kwargs):
@@ -374,7 +376,7 @@ def check_element(driver, id_number):
     checks that an element exists on a website, and provides an exception
     """
     try:
-        driver.find_element_by_id(id_number)
+        driver.find_element('id', id_number)
     except NoSuchElementException as e:
         print(e)
         return False
@@ -549,7 +551,10 @@ def frequency_to_bandname(frequency):
 
 def fetch_driver():
     # open the webdriver
-    return webdriver.PhantomJS()
+    options = webdriver.ChromeOptions()
+    options.add_argument("--headless=new")
+    driver = webdriver.Chrome(options=options)
+    return driver
 
 
 def calc_credible_intervals(samples, interval=0.9):
@@ -785,6 +790,265 @@ def interpolated_barnes_and_kasen_thermalisation_efficiency(mej, vej):
     return av, bv, dv
 
 
+def heatinggrids():
+    # Grid of velocity and Ye
+    YE_GRID = np.array([0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5], dtype=np.float64)
+    V_GRID = np.array([0.05, 0.1, 0.2, 0.3, 0.4, 0.5], dtype=np.float64)
+
+    # Approximant coefficients on the grid
+    E0_GRID = np.array([
+        1.000, 1.000, 1.000, 1.000, 1.000, 1.000,
+        1.000, 1.000, 1.041, 1.041, 1.041, 1.041,
+        1.146, 1.000, 1.041, 1.041, 1.041, 1.041,
+        1.146, 1.000, 1.000, 1.000, 1.041, 1.041,
+        1.301, 1.398, 1.602, 1.580, 1.763, 1.845,
+        0.785, 1.255, 1.673, 1.673, 1.874, 1.874,
+        0.863, 0.845, 1.212, 1.365, 1.635, 2.176,
+        -2.495, -2.495, -2.097, -2.155, -2.046, -1.824,
+        -0.699, -0.699, -0.222, 0.176, 0.176, 0.176,
+        -0.398, 0.000, 0.301, 0.477, 0.477, 0.477], dtype=np.float64)
+
+    # Reshape GRIDs to a 2D array
+    E0_GRID = E0_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # ALP_GRID
+    ALP_GRID = np.array([
+        1.37, 1.38, 1.41, 1.41, 1.41, 1.41,
+        1.41, 1.38, 1.37, 1.37, 1.37, 1.37,
+        1.41, 1.38, 1.37, 1.37, 1.37, 1.37,
+        1.36, 1.25, 1.32, 1.32, 1.34, 1.34,
+        1.44, 1.40, 1.46, 1.66, 1.60, 1.60,
+        1.36, 1.33, 1.33, 1.33, 1.374, 1.374,
+        1.40, 1.358, 1.384, 1.384, 1.384, 1.344,
+        1.80, 1.80, 2.10, 2.10, 1.90, 1.90,
+        8.00, 8.00, 7.00, 7.00, 7.00, 7.00,
+        1.40, 1.40, 1.40, 1.60, 1.60, 1.60
+    ], dtype=np.float64)
+
+    ALP_GRID = ALP_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # T0_GRID
+    T0_GRID = np.array([
+        1.80, 1.40, 1.20, 1.20, 1.20, 1.20,
+        1.40, 1.00, 0.85, 0.85, 0.85, 0.85,
+        1.00, 0.80, 0.65, 0.65, 0.61, 0.61,
+        0.85, 0.60, 0.45, 0.45, 0.45, 0.45,
+        0.65, 0.38, 0.22, 0.18, 0.12, 0.095,
+        0.540, 0.31, 0.18, 0.13, 0.095, 0.081,
+        0.385, 0.235, 0.1, 0.06, 0.035, 0.025,
+        26.0, 26.0, 0.4, 0.4, 0.12, -20.0,
+        0.20, 0.12, 0.05, 0.03, 0.025, 0.021,
+        0.16, 0.08, 0.04, 0.02, 0.018, 0.016
+    ], dtype=np.float64)
+
+    T0_GRID = T0_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # SIG_GRID
+    SIG_GRID = np.array([
+        0.08, 0.08, 0.095, 0.095, 0.095, 0.095,
+        0.10, 0.08, 0.070, 0.070, 0.070, 0.070,
+        0.07, 0.08, 0.070, 0.065, 0.070, 0.070,
+        0.040, 0.030, 0.05, 0.05, 0.05, 0.050,
+        0.05, 0.030, 0.025, 0.045, 0.05, 0.05,
+        0.11, 0.04, 0.021, 0.021, 0.017, 0.017,
+        0.10, 0.094, 0.068, 0.05, 0.03, 0.01,
+        45.0, 45.0, 45.0, 45.0, 25.0, 40.0,
+        0.20, 0.12, 0.05, 0.03, 0.025, 0.021,
+        0.03, 0.015, 0.007, 0.01, 0.009, 0.007
+    ], dtype=np.float64)
+
+    SIG_GRID = SIG_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # ALP1_GRID
+    ALP1_GRID = np.array([
+        7.50, 7.50, 7.50, 7.50, 7.50, 7.50,
+        9.00, 9.00, 7.50, 7.50, 7.00, 7.00,
+        8.00, 8.00, 7.50, 7.50, 7.00, 7.00,
+        8.00, 8.00, 7.50, 7.50, 7.00, 7.00,
+        8.00, 8.00, 5.00, 7.50, 7.00, 6.50,
+        4.5, 3.8, 4.0, 4.0, 4.0, 4.0,
+        2.4, 3.8, 3.8, 3.21, 2.91, 3.61,
+        -1.55, -1.55, -0.75, -0.75, -2.50, -5.00,
+        -1.55, -1.55, -1.55, -1.55, -1.55, -1.55,
+        3.00, 3.00, 3.00, 3.00, 3.00, 3.00
+    ], dtype=np.float64)
+
+    ALP1_GRID = ALP1_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # T1_GRID
+    T1_GRID = np.array([
+        0.040, 0.025, 0.014, 0.010, 0.008, 0.006,
+        0.040, 0.035, 0.020, 0.012, 0.010, 0.008,
+        0.080, 0.040, 0.020, 0.012, 0.012, 0.009,
+        0.080, 0.040, 0.030, 0.018, 0.012, 0.009,
+        0.080, 0.060, 0.065, 0.028, 0.020, 0.015,
+        0.14, 0.123, 0.089, 0.060, 0.045, 0.031,
+        0.264, 0.1, 0.07, 0.055, 0.042, 0.033,
+        1.0, 1.0, 1.0, 1.0, 0.02, 0.01,
+        1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
+        0.04, 0.02, 0.01, 0.002, 0.002, 0.002
+    ], dtype=np.float64)
+
+    T1_GRID = T1_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    SIG1_GRID = np.array([0.250, 0.120, 0.045, 0.028, 0.020, 0.015,
+                          0.250, 0.060, 0.035, 0.020, 0.016, 0.012,
+                          0.170, 0.090, 0.035, 0.020, 0.012, 0.009,
+                          0.170, 0.070, 0.035, 0.015, 0.012, 0.009,
+                          0.170, 0.070, 0.050, 0.025, 0.020, 0.020,
+                          0.065, 0.067, 0.053, 0.032, 0.032, 0.024,
+                          0.075, 0.044, 0.03, 0.02, 0.02, 0.014,
+                          10.0, 10.0, 10.0, 10.0, 0.02, 0.01,
+                          10.0, 10.0, 10.0, 10.0, 10.0, 10.0,
+                          0.01, 0.005, 0.002, 1e-4, 1e-4, 1e-4])
+
+    SIG1_GRID = SIG1_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    C1_GRID = np.array([27.2, 27.8, 28.2, 28.2, 28.2, 28.2,
+                        28.0, 27.8, 27.8, 27.8, 27.8, 27.8,
+                        27.5, 27.0, 27.8, 27.8, 27.8, 27.8,
+                        28.8, 28.1, 27.8, 27.8, 27.5, 27.5,
+                        28.5, 28.0, 27.5, 28.5, 29.2, 29.0,
+                        25.0, 27.5, 25.8, 20.9, 29.3, 1.0,
+                        28.7, 27.0, 28.0, 28.0, 27.4, 25.3,
+                        28.5, 29.1, 29.5, 30.1, 30.4, 29.9,
+                        20.4, 20.6, 20.8, 20.9, 20.9, 21.0,
+                        29.9, 30.1, 30.1, 30.2, 30.3, 30.3])
+
+    C1_GRID = C1_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    TAU1_GRID = np.array([4.07, 4.07, 4.07, 4.07, 4.07, 4.07,
+                          4.07, 4.07, 4.07, 4.07, 4.07, 4.07,
+                          4.07, 4.07, 4.07, 4.07, 4.07, 4.07,
+                          4.07, 4.07, 4.07, 4.07, 4.07, 4.07,
+                          4.77, 4.77, 4.77, 4.77, 4.07, 4.07,
+                          4.77, 4.77, 28.2, 1.03, 0.613, 1.0,
+                          3.4, 14.5, 11.4, 14.3, 13.3, 13.3,
+                          2.52, 2.52, 2.52, 2.52, 2.52, 2.52,
+                          1.02, 1.02, 1.02, 1.02, 1.02, 1.02,
+                          0.22, 0.22, 0.22, 0.22, 0.22, 0.22])
+
+    TAU1_GRID = TAU1_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    C2_GRID = np.array([21.5, 21.5, 22.1, 22.1, 22.1, 22.1,
+                        22.3, 21.5, 21.5, 21.8, 21.8, 21.8,
+                        22.0, 21.5, 21.5, 22.0, 21.8, 21.8,
+                        23.5, 22.5, 22.1, 22.0, 22.2, 22.2,
+                        22.0, 22.8, 23.0, 23.0, 23.5, 23.5,
+                        10.0, 0.0, 0.0, 19.8, 22.0, 21.0,
+                        26.2, 14.1, 18.8, 19.1, 23.8, 19.2,
+                        25.4, 25.4, 25.8, 26.0, 26.0, 25.8,
+                        18.4, 18.4, 18.6, 18.6, 18.6, 18.6,
+                        27.8, 28.0, 28.2, 28.2, 28.3, 28.3])
+
+    C2_GRID = C2_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    TAU2_GRID = np.array([4.62, 4.62, 4.62, 4.62, 4.62, 4.62,
+                          4.62, 4.62, 4.62, 4.62, 4.62, 4.62,
+                          4.62, 4.62, 4.62, 4.62, 4.62, 4.62,
+                          4.62, 4.62, 4.62, 4.62, 4.62, 4.62,
+                          5.62, 5.62, 5.62, 5.62, 4.62, 4.62,
+                          5.62, 5.18, 5.18, 34.7, 8.38, 22.6,
+                          0.15, 4.49, 95.0, 95.0, 0.95, 146.,
+                          0.12, 0.12, 0.12, 0.12, 0.12, 0.14,
+                          0.32, 0.32, 0.32, 0.32, 0.32, 0.32,
+                          0.02, 0.02, 0.02, 0.02, 0.02, 0.02])
+
+    TAU2_GRID = TAU2_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    C3_GRID = np.array([19.4, 19.8, 20.1, 20.1, 20.1, 20.1,
+                        20.0, 19.8, 19.8, 19.8, 19.8, 19.8,
+                        19.9, 19.8, 19.8, 19.8, 19.8, 19.8,
+                        5.9, 9.8, 23.5, 23.5, 23.5, 23.5,
+                        27.3, 26.9, 26.6, 27.4, 25.8, 25.8,
+                        27.8, 26.9, 18.9, 25.4, 24.8, 25.8,
+                        22.8, 17.9, 18.9, 25.4, 24.8, 25.5,
+                        20.6, 20.2, 19.8, 19.2, 19.5, 18.4,
+                        12.6, 13.1, 14.1, 14.5, 14.5, 14.5,
+                        24.3, 24.2, 24.0, 24.0, 24.0, 23.9])
+
+    C3_GRID = C3_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    TAU3_GRID = np.array([18.2, 18.2, 18.2, 18.2, 18.2, 18.2,
+                          18.2, 18.2, 18.2, 18.2, 18.2, 18.2,
+                          18.2, 18.2, 18.2, 18.2, 18.2, 18.2,
+                          18.2, 18.2, 0.62, 0.62, 0.62, 0.62,
+                          0.18, 0.18, 0.18, 0.18, 0.32, 0.32,
+                          0.12, 0.18, 50.8, 0.18, 0.32, 0.32,
+                          2.4, 51.8, 50.8, 0.18, 0.32, 0.32,
+                          3.0, 2.5, 2.4, 2.4, 2.4, 60.4,
+                          200., 200., 200., 200., 200., 200.,
+                          8.76, 8.76, 8.76, 8.76, 8.76, 8.76])
+
+    TAU3_GRID = TAU3_GRID.reshape((len(V_GRID), len(YE_GRID)), order='F')
+
+    # make interpolants
+    E0_interp = RegularGridInterpolator((V_GRID, YE_GRID), E0_GRID, bounds_error=False, fill_value=None)
+    ALP_interp = RegularGridInterpolator((V_GRID, YE_GRID), ALP_GRID, bounds_error=False, fill_value=None)
+    T0_interp = RegularGridInterpolator((V_GRID, YE_GRID), T0_GRID, bounds_error=False, fill_value=None)
+    SIG_interp = RegularGridInterpolator((V_GRID, YE_GRID), SIG_GRID, bounds_error=False, fill_value=None)
+    ALP1_interp = RegularGridInterpolator((V_GRID, YE_GRID), ALP1_GRID, bounds_error=False, fill_value=None)
+    T1_interp = RegularGridInterpolator((V_GRID, YE_GRID), T1_GRID, bounds_error=False, fill_value=None)
+    SIG1_interp = RegularGridInterpolator((V_GRID, YE_GRID), SIG1_GRID, bounds_error=False, fill_value=None)
+    C1_interp = RegularGridInterpolator((V_GRID, YE_GRID), C1_GRID, bounds_error=False, fill_value=None)
+    TAU1_interp = RegularGridInterpolator((V_GRID, YE_GRID), TAU1_GRID, bounds_error=False, fill_value=None)
+    C2_interp = RegularGridInterpolator((V_GRID, YE_GRID), C2_GRID, bounds_error=False, fill_value=None)
+    TAU2_interp = RegularGridInterpolator((V_GRID, YE_GRID), TAU2_GRID, bounds_error=False, fill_value=None)
+    C3_interp = RegularGridInterpolator((V_GRID, YE_GRID), C3_GRID, bounds_error=False, fill_value=None)
+    TAU3_interp = RegularGridInterpolator((V_GRID, YE_GRID), TAU3_GRID, bounds_error=False, fill_value=None)
+
+
+    interpolators = namedtuple('interpolators', ['E0', 'ALP', 'T0', 'SIG', 'ALP1', 'T1', 'SIG1',
+                                                 'C1', 'TAU1', 'C2', 'TAU2', 'C3', 'TAU3'])
+    interpolators.E0 = E0_interp
+    interpolators.ALP = ALP_interp
+    interpolators.T0 = T0_interp
+    interpolators.SIG = SIG_interp
+    interpolators.ALP1 = ALP1_interp
+    interpolators.T1 = T1_interp
+    interpolators.SIG1 = SIG1_interp
+    interpolators.C1 = C1_interp
+    interpolators.TAU1 = TAU1_interp
+    interpolators.C2 = C2_interp
+    interpolators.TAU2 = TAU2_interp
+    interpolators.C3 = C3_interp
+    interpolators.TAU3 = TAU3_interp
+    return interpolators
+
+def get_heating_terms(ye, vel, **kwargs):
+    ints = heatinggrids()
+    e0 = ints.E0([vel, ye])[0]
+    alp = ints.ALP([vel, ye])[0]
+    t0 = ints.T0([vel, ye])[0]
+    sig = ints.SIG([vel, ye])[0]
+    alp1 = ints.ALP1([vel, ye])[0]
+    t1 = ints.T1([vel, ye])[0]
+    sig1 = ints.SIG1([vel, ye])[0]
+    c1 = ints.C1([vel, ye])[0]
+    tau1 = ints.TAU1([vel, ye])[0]
+    c2 = ints.C2([vel, ye])[0]
+    tau2 = ints.TAU2([vel, ye])[0]
+    c3 = ints.C3([vel, ye])[0]
+    tau3 = ints.TAU3([vel, ye])[0]
+    heating_terms = namedtuple('heating_terms', ['e0', 'alp', 't0', 'sig', 'alp1', 't1', 'sig1', 'c1',
+                                                 'tau1', 'c2', 'tau2', 'c3', 'tau3'])
+
+    heating_rate_fudge = kwargs.get('heating_rate_fudge', 1.0)
+    heating_terms.e0 = e0 * heating_rate_fudge
+    heating_terms.alp = alp * heating_rate_fudge
+    heating_terms.t0 = t0 * heating_rate_fudge
+    heating_terms.sig = sig * heating_rate_fudge
+    heating_terms.alp1 = alp1 * heating_rate_fudge
+    heating_terms.t1 = t1 * heating_rate_fudge
+    heating_terms.sig1 = sig1 * heating_rate_fudge
+    heating_terms.c1 = c1 * heating_rate_fudge
+    heating_terms.tau1 = tau1 * heating_rate_fudge
+    heating_terms.c2 = c2 * heating_rate_fudge
+    heating_terms.tau2 = tau2 * heating_rate_fudge
+    heating_terms.c3 = c3 * heating_rate_fudge
+    heating_terms.tau3 = tau3 * heating_rate_fudge
+    return heating_terms
+
 def electron_fraction_from_kappa(kappa):
     """
     Uses interpolation from Tanaka+19 to calculate
@@ -795,9 +1059,22 @@ def electron_fraction_from_kappa(kappa):
 
     kappa_array = np.array([1, 3, 5, 20, 30])
     ye_array = np.array([0.4,0.35,0.25,0.2, 0.1])
-    kappa_func = interp1d(kappa_array, y=ye_array)
+    kappa_func = interp1d(kappa_array, y=ye_array, fill_value='extrapolate')
     electron_fraction = kappa_func(kappa)
     return electron_fraction
+
+def kappa_from_electron_fraction(ye):
+    """
+    Uses interpolation from Tanaka+19 to calculate
+    the opacity based on the electron fraction
+    :param ye: electron fraction
+    :return: electron_fraction
+    """
+    kappa_array = np.array([1, 3, 5, 20, 30])
+    ye_array = np.array([0.4,0.35,0.25,0.2, 0.1])
+    func = interp1d(ye_array, y=kappa_array, fill_value='extrapolate')
+    kappa = func(ye)
+    return kappa
     
 def lorentz_factor_from_velocity(velocity):
     """