From 1766868e80a8c292a1b6c91650b62cf95e38b620 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Fri, 23 Feb 2024 16:52:11 +0100 Subject: [PATCH 01/27] Implement new heating rate kilonova model following Rosswog and Korokbin 2022. --- redback/transient_models/kilonova_models.py | 202 ++++++++++++++- redback/utils.py | 270 ++++++++++++++++++++ 2 files changed, 471 insertions(+), 1 deletion(-) diff --git a/redback/transient_models/kilonova_models.py b/redback/transient_models/kilonova_models.py index 72ec5a1e..196fe7c5 100644 --- a/redback/transient_models/kilonova_models.py +++ b/redback/transient_models/kilonova_models.py @@ -10,7 +10,7 @@ from redback.interaction_processes import Diffusion, AsphericalDiffusion from redback.utils import calc_kcorrected_properties, interpolated_barnes_and_kasen_thermalisation_efficiency, \ - electron_fraction_from_kappa, citation_wrapper, lambda_to_nu + electron_fraction_from_kappa, citation_wrapper, lambda_to_nu, get_heating_terms, kappa_from_electron_fraction from redback.eos import PiecewisePolytrope from redback.sed import blackbody_to_flux_density, get_correct_output_format_from_spectra, Blackbody from redback.constants import * @@ -1421,6 +1421,206 @@ def one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs): spectra=spectra, lambda_array=lambda_observer_frame, **kwargs) + +def _calc_new_heating_rate(time, mej, electron_fraction, ejecta_velocity): + heating_terms = get_heating_terms(electron_fraction, ejecta_velocity) + # rescale + m0 = mej * solar_mass + c1 = np.exp(heating_terms.c1) + c2 = np.exp(heating_terms.c2) + c3 = np.exp(heating_terms.c3) + tau1 = 1e3*heating_terms.tau1 + tau2 = 1e5*heating_terms.tau2 + tau3 = 1e5*heating_terms.tau3 + term1 = heating_terms.e0 * 1e18 * (0.5 - np.arctan((time - heating_terms.t0) / heating_terms.sig) / np.pi)**heating_terms.alp + term2 = (0.5 + np.arctan((time - heating_terms.t1)/heating_terms.sig1) / np.pi )**heating_terms.alp1 + term3 = c1 * np.exp(-time/tau1) + term4 = c2 * np.exp(-time/tau2) + term5 = c3 * np.exp(-time/tau3) + lum_in = term1*term2 + term3 + term4 + term5 + return lum_in*m0 + +def _one_component_kilonova_rosswog_heatingrate(time, mej, vej, electron_fraction, **kwargs): + tdays = time/day_to_s + # set up kilonova physics + av, bv, dv = interpolated_barnes_and_kasen_thermalisation_efficiency(mej, vej) + # thermalisation from Barnes+16 + e_th = 0.36 * (np.exp(-av * tdays) + np.log1p(2.0 * bv * tdays ** dv) / (2.0 * bv * tdays ** dv)) + temperature_floor = kwargs.get('temperature_floor', 4000) # kelvin + beta = 13.7 + + v0 = vej * speed_of_light + m0 = mej * solar_mass + kappa = kappa_from_electron_fraction(electron_fraction) + tdiff = np.sqrt(2.0 * kappa * (m0) / (beta * v0 * speed_of_light)) + + lum_in = _calc_new_heating_rate(time, mej, electron_fraction, vej) + integrand = lum_in * e_th * (time / tdiff) * np.exp(time ** 2 / tdiff ** 2) + + bolometric_luminosity = np.zeros(len(time)) + bolometric_luminosity[1:] = cumtrapz(integrand, time) + bolometric_luminosity[0] = bolometric_luminosity[1] + bolometric_luminosity = bolometric_luminosity * np.exp(-time ** 2 / tdiff ** 2) / tdiff + + temperature = (bolometric_luminosity / (4.0 * np.pi * sigma_sb * v0 ** 2 * time ** 2)) ** 0.25 + r_photosphere = (bolometric_luminosity / (4.0 * np.pi * sigma_sb * temperature_floor ** 4)) ** 0.5 + + # check temperature floor conditions + mask = temperature <= temperature_floor + temperature[mask] = temperature_floor + mask = np.logical_not(mask) + r_photosphere[mask] = v0 * time[mask] + return bolometric_luminosity, temperature, r_photosphere + +@citation_wrapper('redback') +def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, electron_fraction, **kwargs): + """ + :param time: observer frame time in days + :param redshift: redshift + :param mej: ejecta mass in solar masses + :param vej: minimum initial velocity + :param kappa: gray opacity + :param kwargs: Additional keyword arguments + :param temperature_floor: Temperature floor in K (default 4000) + :param frequency: Required if output_format is 'flux_density'. + frequency to calculate - Must be same length as time array or a single number). + :param bands: Required if output_format is 'magnitude' or 'flux'. + :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source' + :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on. + :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object. + :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source' + """ + cosmology = kwargs.get('cosmology', cosmo) + dl = cosmology.luminosity_distance(redshift).cgs.value + time_temp = np.geomspace(1e-3, 5e6, 300) # in source frame + time_obs = time + _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej, vej, electron_fraction, **kwargs) + + if kwargs['output_format'] == 'flux_density': + time = time_obs * day_to_s + frequency = kwargs['frequency'] + # interpolate properties onto observation times + temp_func = interp1d(time_temp, y=temperature) + rad_func = interp1d(time_temp, y=r_photosphere) + # convert to source frame time and frequency + frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time) + + temp = temp_func(time) + photosphere = rad_func(time) + + flux_density = blackbody_to_flux_density(temperature=temp, r_photosphere=photosphere, + dl=dl, frequency=frequency) + + return flux_density.to(uu.mJy).value + + else: + lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 200)) + time_observer_frame = time_temp * (1. + redshift) + frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame), + redshift=redshift, time=time_observer_frame) + fmjy = blackbody_to_flux_density(temperature=temperature, + r_photosphere=r_photosphere, frequency=frequency[:, None], dl=dl) + fmjy = fmjy.T + spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom, + equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom)) + if kwargs['output_format'] == 'spectra': + return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame, + lambdas=lambda_observer_frame, + spectra=spectra) + else: + return get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame/day_to_s, + spectra=spectra, lambda_array=lambda_observer_frame, + **kwargs) + +@citation_wrapper('redback') +def two_comp_kne_rosswog_heatingrate(time, redshift, mej_1, vej_1, temperature_floor_1, ye_1, + mej_2, vej_2, temperature_floor_2, ye_2, **kwargs): + """ + :param time: observer frame time in days + :param redshift: redshift + :param mej_1: ejecta mass in solar masses of first component + :param vej_1: minimum initial velocity of first component + :param kappa_1: gray opacity of first component + :param temperature_floor_1: floor temperature of first component + :param mej_2: ejecta mass in solar masses of second component + :param vej_2: minimum initial velocity of second component + :param temperature_floor_2: floor temperature of second component + :param kappa_2: gray opacity of second component + :param kwargs: Additional keyword arguments + :param frequency: Required if output_format is 'flux_density'. + frequency to calculate - Must be same length as time array or a single number). + :param bands: Required if output_format is 'magnitude' or 'flux'. + :param output_format: 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source' + :param lambda_array: Optional argument to set your desired wavelength array (in Angstroms) to evaluate the SED on. + :param cosmology: Cosmology to use for luminosity distance calculation. Defaults to Planck18. Must be a astropy.cosmology object. + :return: set by output format - 'flux_density', 'magnitude', 'spectra', 'flux', 'sncosmo_source' + """ + cosmology = kwargs.get('cosmology', cosmo) + dl = cosmology.luminosity_distance(redshift).cgs.value + time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame + time_obs = time + + mej = [mej_1, mej_2] + vej = [vej_1, vej_2] + temperature_floor = [temperature_floor_1, temperature_floor_2] + ye = [ye_1, ye_2] + + if kwargs['output_format'] == 'flux_density': + time = time * day_to_s + frequency = kwargs['frequency'] + + # convert to source frame time and frequency + frequency, time = calc_kcorrected_properties(frequency=frequency, redshift=redshift, time=time) + + ff = np.zeros(len(time)) + for x in range(2): + temp_kwargs = {} + temp_kwargs['temperature_floor'] = temperature_floor[x] + _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej[x], vej[x], ye[x], + **temp_kwargs) + # interpolate properties onto observation times + temp_func = interp1d(time_temp, y=temperature) + rad_func = interp1d(time_temp, y=r_photosphere) + temp = temp_func(time) + photosphere = rad_func(time) + flux_density = blackbody_to_flux_density(temperature=temp, r_photosphere=photosphere, + dl=dl, frequency=frequency) + units = flux_density.unit + ff += flux_density.value + + ff = ff * units + return ff.to(uu.mJy).value + + else: + lambda_observer_frame = kwargs.get('lambda_array', np.geomspace(100, 60000, 200)) + time_observer_frame = time_temp * (1. + redshift) + frequency, time = calc_kcorrected_properties(frequency=lambda_to_nu(lambda_observer_frame), + redshift=redshift, time=time_observer_frame) + full_spec = np.zeros((len(time), len(frequency))) + + for x in range(2): + temp_kwargs = {} + temp_kwargs['temperature_floor'] = temperature_floor[x] + _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej[x], vej[x], ye[x], + **temp_kwargs) + fmjy = blackbody_to_flux_density(temperature=temperature, + r_photosphere=r_photosphere, frequency=frequency[:, None], dl=dl) + fmjy = fmjy.T + spectra = fmjy.to(uu.mJy).to(uu.erg / uu.cm ** 2 / uu.s / uu.Angstrom, + equivalencies=uu.spectral_density(wav=lambda_observer_frame * uu.Angstrom)) + units = spectra.unit + full_spec += spectra.value + + full_spec = full_spec * units + if kwargs['output_format'] == 'spectra': + return namedtuple('output', ['time', 'lambdas', 'spectra'])(time=time_observer_frame, + lambdas=lambda_observer_frame, + spectra=full_spec) + else: + return get_correct_output_format_from_spectra(time=time_obs, time_eval=time_observer_frame/day_to_s, + spectra=full_spec, lambda_array=lambda_observer_frame, + **kwargs) + def _one_component_kilonova_model(time, mej, vej, kappa, **kwargs): """ :param time: source frame time in seconds diff --git a/redback/utils.py b/redback/utils.py index c42a088a..15efda7b 100644 --- a/redback/utils.py +++ b/redback/utils.py @@ -785,6 +785,262 @@ 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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_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(YE_GRID), len(V_GRID)), order='F') + + # make interpolants + E0_interp = RegularGridInterpolator((YE_GRID, V_GRID), E0_GRID, bounds_error=False, fill_value=None) + ALP_interp = RegularGridInterpolator((YE_GRID, V_GRID), ALP_GRID, bounds_error=False, fill_value=None) + T0_interp = RegularGridInterpolator((YE_GRID, V_GRID), T0_GRID, bounds_error=False, fill_value=None) + SIG_interp = RegularGridInterpolator((YE_GRID, V_GRID), SIG_GRID, bounds_error=False, fill_value=None) + ALP1_interp = RegularGridInterpolator((YE_GRID, V_GRID), ALP1_GRID, bounds_error=False, fill_value=None) + T1_interp = RegularGridInterpolator((YE_GRID, V_GRID), T1_GRID, bounds_error=False, fill_value=None) + SIG1_interp = RegularGridInterpolator((YE_GRID, V_GRID), SIG1_GRID, bounds_error=False, fill_value=None) + C1_interp = RegularGridInterpolator((YE_GRID, V_GRID), C1_GRID, bounds_error=False, fill_value=None) + TAU1_interp = RegularGridInterpolator((YE_GRID, V_GRID), TAU1_GRID, bounds_error=False, fill_value=None) + C2_interp = RegularGridInterpolator((YE_GRID, V_GRID), C2_GRID, bounds_error=False, fill_value=None) + TAU2_interp = RegularGridInterpolator((YE_GRID, V_GRID), TAU2_GRID, bounds_error=False, fill_value=None) + C3_interp = RegularGridInterpolator((YE_GRID, V_GRID), C3_GRID, bounds_error=False, fill_value=None) + TAU3_interp = RegularGridInterpolator((YE_GRID, V_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): + ints = heatinggrids() + e0 = ints.E0([ye, vel])[0] + alp = ints.ALP([ye, vel])[0] + t0 = ints.T0([ye, vel])[0] + sig = ints.SIG([ye, vel])[0] + alp1 = ints.ALP1([ye, vel])[0] + t1 = ints.T1([ye, vel])[0] + sig1 = ints.SIG1([ye, vel])[0] + c1 = ints.C1([ye, vel])[0] + tau1 = ints.TAU1([ye, vel])[0] + c2 = ints.C2([ye, vel])[0] + tau2 = ints.TAU2([ye, vel])[0] + c3 = ints.C3([ye, vel])[0] + tau3 = ints.TAU3([ye, vel])[0] + heating_terms = namedtuple('heating_terms', ['e0', 'alp', 't0', 'sig', 'alp1', 't1', 'sig1', 'c1', + 'tau1', 'c2', 'tau2', 'c3', 'tau3']) + heating_terms.e0 = e0 + heating_terms.alp = alp + heating_terms.t0 = t0 + heating_terms.sig = sig + heating_terms.alp1 = alp1 + heating_terms.t1 = t1 + heating_terms.sig1 = sig1 + heating_terms.c1 = c1 + heating_terms.tau1 = tau1 + heating_terms.c2 = c2 + heating_terms.tau2 = tau2 + heating_terms.c3 = c3 + heating_terms.tau3 = tau3 + return heating_terms + def electron_fraction_from_kappa(kappa): """ Uses interpolation from Tanaka+19 to calculate @@ -798,6 +1054,20 @@ def electron_fraction_from_kappa(kappa): kappa_func = interp1d(kappa_array, y=ye_array) 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]) + kappa_func = interp1d(kappa_array, y=ye_array) + electron_fraction = kappa_func(kappa) + return electron_fraction def lorentz_factor_from_velocity(velocity): """ From 916c0718086aaa1f0e629f1da9e30c7eb962c5ef Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Fri, 23 Feb 2024 16:57:14 +0100 Subject: [PATCH 02/27] Improve implementation --- redback/priors/one_comp_kne_rosswog_heatingrate.prior | 5 +++++ redback/priors/two_comp_kne_rosswog_heatingrate.prior | 9 +++++++++ redback/transient_models/kilonova_models.py | 4 ++-- redback/utils.py | 6 +++--- 4 files changed, 19 insertions(+), 5 deletions(-) create mode 100644 redback/priors/one_comp_kne_rosswog_heatingrate.prior create mode 100644 redback/priors/two_comp_kne_rosswog_heatingrate.prior diff --git a/redback/priors/one_comp_kne_rosswog_heatingrate.prior b/redback/priors/one_comp_kne_rosswog_heatingrate.prior new file mode 100644 index 00000000..2d37cedd --- /dev/null +++ b/redback/priors/one_comp_kne_rosswog_heatingrate.prior @@ -0,0 +1,5 @@ +redshift = Uniform(1e-6, 0.1, 'redshift', latex_label = r'$z$') +mej = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') +vej = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +temperature_floor = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') \ No newline at end of file diff --git a/redback/priors/two_comp_kne_rosswog_heatingrate.prior b/redback/priors/two_comp_kne_rosswog_heatingrate.prior new file mode 100644 index 00000000..4b43042d --- /dev/null +++ b/redback/priors/two_comp_kne_rosswog_heatingrate.prior @@ -0,0 +1,9 @@ +redshift = Uniform(1e-6, 0.1, 'redshift', latex_label = r'$z$') +mej_1 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') +vej_1 = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye_1 = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +temperature_floor_1 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') +mej_2 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') +vej_2 = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye_2 = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +temperature_floor_2 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') \ No newline at end of file diff --git a/redback/transient_models/kilonova_models.py b/redback/transient_models/kilonova_models.py index 196fe7c5..41fca7c2 100644 --- a/redback/transient_models/kilonova_models.py +++ b/redback/transient_models/kilonova_models.py @@ -1473,7 +1473,7 @@ def _one_component_kilonova_rosswog_heatingrate(time, mej, vej, electron_fractio return bolometric_luminosity, temperature, r_photosphere @citation_wrapper('redback') -def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, electron_fraction, **kwargs): +def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, ye, **kwargs): """ :param time: observer frame time in days :param redshift: redshift @@ -1494,7 +1494,7 @@ def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, electron_fraction dl = cosmology.luminosity_distance(redshift).cgs.value time_temp = np.geomspace(1e-3, 5e6, 300) # in source frame time_obs = time - _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej, vej, electron_fraction, **kwargs) + _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej, vej, ye, **kwargs) if kwargs['output_format'] == 'flux_density': time = time_obs * day_to_s diff --git a/redback/utils.py b/redback/utils.py index 15efda7b..01a01981 100644 --- a/redback/utils.py +++ b/redback/utils.py @@ -1065,9 +1065,9 @@ def kappa_from_electron_fraction(ye): 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) - electron_fraction = kappa_func(kappa) - return electron_fraction + func = interp1d(ye_array, y=kappa_array) + kappa = func(ye) + return kappa def lorentz_factor_from_velocity(velocity): """ From 4f669640753e172a36650e726c695ca42387f956 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Fri, 23 Feb 2024 17:00:04 +0100 Subject: [PATCH 03/27] Improve implementation even more --- redback/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/redback/utils.py b/redback/utils.py index 01a01981..a2b6908b 100644 --- a/redback/utils.py +++ b/redback/utils.py @@ -1062,10 +1062,9 @@ def kappa_from_electron_fraction(ye): :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) + func = interp1d(ye_array, y=kappa_array, fill_value='extrapolate') kappa = func(ye) return kappa From 8a6c3bb7924c98c66ad540c50bd973b27682226b Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Fri, 23 Feb 2024 17:07:33 +0100 Subject: [PATCH 04/27] Update prior --- redback/priors/one_comp_kne_rosswog_heatingrate.prior | 4 ++-- redback/priors/two_comp_kne_rosswog_heatingrate.prior | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/redback/priors/one_comp_kne_rosswog_heatingrate.prior b/redback/priors/one_comp_kne_rosswog_heatingrate.prior index 2d37cedd..ddd7761b 100644 --- a/redback/priors/one_comp_kne_rosswog_heatingrate.prior +++ b/redback/priors/one_comp_kne_rosswog_heatingrate.prior @@ -1,5 +1,5 @@ redshift = Uniform(1e-6, 0.1, 'redshift', latex_label = r'$z$') mej = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') -vej = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') -ye = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +vej = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}$') temperature_floor = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') \ No newline at end of file diff --git a/redback/priors/two_comp_kne_rosswog_heatingrate.prior b/redback/priors/two_comp_kne_rosswog_heatingrate.prior index 4b43042d..a1e5f421 100644 --- a/redback/priors/two_comp_kne_rosswog_heatingrate.prior +++ b/redback/priors/two_comp_kne_rosswog_heatingrate.prior @@ -1,9 +1,9 @@ redshift = Uniform(1e-6, 0.1, 'redshift', latex_label = r'$z$') mej_1 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') -vej_1 = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') -ye_1 = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +vej_1 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye_1 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}$') temperature_floor_1 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') mej_2 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') -vej_2 = Uniform(0.05, 0.5, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') -ye_2 = Uniform(0.05, 0.5, 'ye', latex_label = r'$Y_{e}$') +vej_2 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') +ye_2 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}$') temperature_floor_2 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') \ No newline at end of file From 22096019b5aaac6c95d748bf79ceaa890fc098ab Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 28 Feb 2024 11:42:11 +0100 Subject: [PATCH 05/27] Clean up docstring --- redback/transient/afterglow.py | 2 +- redback/transient/kilonova.py | 2 +- redback/transient/supernova.py | 2 +- redback/transient/tde.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/redback/transient/afterglow.py b/redback/transient/afterglow.py index c13e079e..6bd6deb0 100644 --- a/redback/transient/afterglow.py +++ b/redback/transient/afterglow.py @@ -57,7 +57,7 @@ def __init__( :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 diff --git a/redback/transient/kilonova.py b/redback/transient/kilonova.py index dd7e69a3..a832818c 100644 --- a/redback/transient/kilonova.py +++ b/redback/transient/kilonova.py @@ -50,7 +50,7 @@ def __init__( :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 diff --git a/redback/transient/supernova.py b/redback/transient/supernova.py index ffbcdab6..4dd3da3b 100644 --- a/redback/transient/supernova.py +++ b/redback/transient/supernova.py @@ -45,7 +45,7 @@ def __init__( :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 diff --git a/redback/transient/tde.py b/redback/transient/tde.py index 7c12c5b0..f504f92c 100644 --- a/redback/transient/tde.py +++ b/redback/transient/tde.py @@ -42,7 +42,7 @@ def __init__( :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 From c3492418c6d767817f52c30f6ad18dc05208a1a4 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sat, 2 Mar 2024 16:57:04 +0100 Subject: [PATCH 06/27] Fix syntax for saving generic simulated transient --- redback/simulate_transients.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redback/simulate_transients.py b/redback/simulate_transients.py index 81d12a5e..fc95e1b5 100644 --- a/redback/simulate_transients.py +++ b/redback/simulate_transients.py @@ -98,7 +98,7 @@ def save_transient(self, name): 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=pd.DataFrame.from_dict([self.parameters]) self.parameters.to_csv(injection_path, index=False) class SimulateOpticalTransient(object): From 4d46c31d7b91bd840b3fc5414ef14437260a5433 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sat, 2 Mar 2024 18:08:27 +0100 Subject: [PATCH 07/27] Add a line spectrum phenomenological model --- redback/transient_models/phenomenological_models.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/redback/transient_models/phenomenological_models.py b/redback/transient_models/phenomenological_models.py index 283dae64..de2cb818 100644 --- a/redback/transient_models/phenomenological_models.py +++ b/redback/transient_models/phenomenological_models.py @@ -1,5 +1,18 @@ import numpy as np +def line_spectrum(wavelength, line_amp, cont_amp, x0): + """ + A gaussian to add or subtract from a continuum spectrum to mimic absorption or emission lines + + :param wavelength: wavelength array in whatever units + :param line_amp: line amplitude scale + :param cont_amp: Continuum amplitude scale + :param x0: Position of emission line + :return: spectrum in whatever units set by line_amp + """ + spectrum = line_amp / cont_amp * np.exp(-(wavelength - x0) ** 2. / (2 * cont_amp ** 2) ) + return spectrum + def gaussian_rise(time, a_1, peak_time, sigma_t): """ :param time: time array in whatver time units From 787cc58c1b2ff44bb930d08b7fb8ea276ecc8762 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 18:47:50 +0100 Subject: [PATCH 08/27] add doc on priors --- redback/priors.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/redback/priors.py b/redback/priors.py index 7b4dad18..30c38836 100644 --- a/redback/priors.py +++ b/redback/priors.py @@ -9,6 +9,17 @@ def get_priors(model, times=None, y=None, yerr=None, dt=None, **kwargs): + """ + Get the prior for the given model. If the model is a prompt model, the times, y, and yerr must be provided. + + :param model: String referring to a name of a model implemented in Redback. + :param times: Time array + :param y: Y values, arbitrary units + :param yerr: Error on y values, arbitrary units + :param dt: time interval + :param kwargs: Extra arguments to be passed to the prior function + :return: priors: PriorDict object + """ prompt_prior_functions = dict(gaussian=get_gaussian_priors, skew_gaussian=get_skew_gaussian_priors, skew_exponential=get_skew_exponential_priors, fred=get_fred_priors, fred_extended=get_fred_extended_priors) From 6d8dc55a81dc6a56d6672e9d4077efa5888fb1ad Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 20:35:49 +0100 Subject: [PATCH 09/27] Add maximum likelihood estimation to all likelihoods --- redback/likelihoods.py | 150 ++++++++++++++++++++++++++++++++++++----- 1 file changed, 133 insertions(+), 17 deletions(-) diff --git a/redback/likelihoods.py b/redback/likelihoods.py index be17e23b..cd826515 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -3,11 +3,14 @@ import bilby from scipy.special import gammaln +from redback.utils import logger +from bilby.core.prior import DeltaFunction, Constraint class _RedbackLikelihood(bilby.Likelihood): - def __init__(self, x: np.ndarray, y: np.ndarray, function: callable, kwargs: dict = None) -> None: + def __init__(self, x: np.ndarray, y: np.ndarray, function: callable, kwargs: dict = None, priors=None, + fiducial_parameters=None) -> None: """ :param x: The x values. @@ -18,11 +21,19 @@ def __init__(self, x: np.ndarray, y: np.ndarray, function: callable, kwargs: dic :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: Union[dict, None] + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] """ self.x = x self.y = y self.function = function self.kwargs = kwargs + self.priors = priors + self.fiducial_parameters = fiducial_parameters parameters = bilby.core.utils.introspection.infer_parameters_from_function(func=function) super().__init__(parameters=dict.fromkeys(parameters)) @@ -46,11 +57,72 @@ def n(self) -> int: """ return len(self.x) + @property + def parameters_to_be_updated(self): + if self.priors is None: + return None + else: + parameters_to_be_updated = [key for key in self.priors if not isinstance( + self.priors[key], (DeltaFunction, Constraint, float, int))] + return parameters_to_be_updated + + def get_parameter_dictionary_from_list(self, parameter_list): + parameter_dictionary = dict(zip(self.parameters_to_be_updated, parameter_list)) + excluded_parameter_keys = set(self.fiducial_parameters) - set(self.parameters_to_be_updated) + for key in excluded_parameter_keys: + parameter_dictionary[key] = self.fiducial_parameters[key] + return parameter_dictionary + + def get_parameter_list_from_dictionary(self, parameter_dict): + return [parameter_dict[k] for k in self.parameters_to_be_updated] + + def get_bounds_from_priors(self, priors): + bounds = [] + for key in self.parameters_to_be_updated: + bounds.append([priors[key].minimum, priors[key].maximum]) + return bounds + + def lnlike_scipy_maximize(self, parameter_list): + self.parameters.update(self.get_parameter_dictionary_from_list(parameter_list)) + return -self.log_likelihood() + + def find_maximum_likelihood_parameters(self, iterations=5, maximization_kwargs=None): + from scipy.optimize import differential_evolution + parameter_bounds = self.get_bounds_from_priors(self.priors) + if self.priors is None: + raise ValueError("Priors must be provided to use this functionality") + if maximization_kwargs is None: + maximization_kwargs = dict() + self.parameters.update(self.fiducial_parameters) + self.parameters["fiducial"] = 0 + updated_parameters_list = self.get_parameter_list_from_dictionary(self.fiducial_parameters) + old_fiducial_ln_likelihood = self.log_likelihood() + for it in range(iterations): + logger.info(f"Optimizing fiducial parameters. Iteration : {it + 1}") + print(f"Optimizing fiducial parameters. Iteration : {it + 1}") + output = differential_evolution( + self.lnlike_scipy_maximize, + bounds=parameter_bounds, + x0=updated_parameters_list, + **maximization_kwargs, + ) + updated_parameters_list = output['x'] + updated_parameters = self.get_parameter_dictionary_from_list(updated_parameters_list) + self.parameters.update(updated_parameters) + new_fiducial_ln_likelihood = self.log_likelihood_ratio() + logger.info(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") + print(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") + if new_fiducial_ln_likelihood - old_fiducial_ln_likelihood < 0.1: + break + old_fiducial_ln_likelihood = new_fiducial_ln_likelihood + return updated_parameters + class GaussianLikelihood(_RedbackLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma: Union[float, None, np.ndarray], - function: callable, kwargs: dict = None) -> None: + function: callable, kwargs: dict = None, priors=None, + fiducial_parameters=None) -> None: """A general Gaussian likelihood - the parameters are inferred from the arguments of function. :param x: The x values. @@ -67,10 +139,17 @@ def __init__( :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: dict + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] """ self._noise_log_likelihood = None - super().__init__(x=x, y=y, function=function, kwargs=kwargs) + super().__init__(x=x, y=y, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) self.sigma = sigma if self.sigma is None: self.parameters['sigma'] = None @@ -121,7 +200,8 @@ def _gaussian_log_likelihood(res: np.ndarray, sigma: Union[float, np.ndarray]) - class GaussianLikelihoodUniformXErrors(GaussianLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma: Union[float, None, np.ndarray], - bin_size: Union[float, None, np.ndarray], function: callable, kwargs: dict = None) -> None: + bin_size: Union[float, None, np.ndarray], function: callable, kwargs: dict = None, priors=None, + fiducial_parameters=None) -> None: """A general Gaussian likelihood with uniform errors in x- the parameters are inferred from the arguments of function. Takes into account the X errors with a Uniform likelihood between the bin high and bin low values. Note that the prior for the true x values must be uniform in this range! @@ -142,9 +222,16 @@ def __init__( :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: dict + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] """ - super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs) + super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) self.xerr = bin_size * np.ones(self.n) def noise_log_likelihood(self) -> float: @@ -183,7 +270,7 @@ def log_likelihood(self) -> float: class GaussianLikelihoodQuadratureNoise(GaussianLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma_i: Union[float, None, np.ndarray], - function: callable, kwargs: dict = None) -> None: + function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """ A general Gaussian likelihood - the parameters are inferred from the arguments of function @@ -205,7 +292,8 @@ def __init__( """ self.sigma_i = sigma_i # These lines of code infer the parameters from the provided function - super().__init__(x=x, y=y, sigma=sigma_i, function=function, kwargs=kwargs) + super().__init__(x=x, y=y, sigma=sigma_i, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) @property def full_sigma(self) -> Union[float, np.ndarray]: @@ -234,7 +322,7 @@ def log_likelihood(self) -> float: class GaussianLikelihoodWithSystematicNoise(GaussianLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma_i: Union[float, None, np.ndarray], - function: callable, kwargs: dict = None) -> None: + function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """ A general Gaussian likelihood - the parameters are inferred from the arguments of function @@ -253,10 +341,17 @@ def __init__( :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: dict + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] """ self.sigma_i = sigma_i # These lines of code infer the parameters from the provided function - super().__init__(x=x, y=y, sigma=sigma_i, function=function, kwargs=kwargs) + super().__init__(x=x, y=y, sigma=sigma_i, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) @property def full_sigma(self) -> Union[float, np.ndarray]: @@ -286,7 +381,7 @@ def log_likelihood(self) -> float: class GaussianLikelihoodQuadratureNoiseNonDetections(GaussianLikelihoodQuadratureNoise): def __init__( self, x: np.ndarray, y: np.ndarray, sigma_i: Union[float, np.ndarray], function: callable, - kwargs: dict = None, upperlimit_kwargs: dict = None) -> None: + kwargs: dict = None, upperlimit_kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """A general Gaussian likelihood - the parameters are inferred from the arguments of function. Takes into account non-detections with a Uniform likelihood for those points @@ -304,8 +399,15 @@ def __init__( :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: dict - """ - super().__init__(x=x, y=y, sigma_i=sigma_i, function=function, kwargs=kwargs) + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] + """ + super().__init__(x=x, y=y, sigma_i=sigma_i, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) self.upperlimit_kwargs = upperlimit_kwargs @property @@ -345,7 +447,7 @@ class GRBGaussianLikelihood(GaussianLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma: Union[float, np.ndarray], - function: callable, kwargs: dict = None) -> None: + function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """A general Gaussian likelihood - the parameters are inferred from the arguments of function. @@ -363,14 +465,21 @@ def __init__( :type function: callable :param kwargs: Any additional keywords for 'function'. :type kwargs: dict + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] """ - super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs) + super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) class PoissonLikelihood(_RedbackLikelihood): def __init__( self, time: np.ndarray, counts: np.ndarray, function: callable, integrated_rate_function: bool = True, - dt: Union[float, np.ndarray] = None, kwargs: dict = None) -> None: + dt: Union[float, np.ndarray] = None, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """ :param time: The time values. :type time: np.ndarray @@ -389,8 +498,15 @@ def __init__( :type dt: Union[float, None, np.ndarray] :param kwargs: Any additional keywords for 'function'. :type kwargs: dict - """ - super(PoissonLikelihood, self).__init__(x=time, y=counts, function=function, kwargs=kwargs) + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] + """ + super(PoissonLikelihood, self).__init__(x=time, y=counts, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) self.integrated_rate_function = integrated_rate_function self.dt = dt self.parameters['background_rate'] = 0 From 43afe753ce581ebee1fd8f09b679a3faed374088 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 20:52:05 +0100 Subject: [PATCH 10/27] unit test for maximum likelihood --- test/likelihood_test.py | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/test/likelihood_test.py b/test/likelihood_test.py index 2bba357b..961377aa 100644 --- a/test/likelihood_test.py +++ b/test/likelihood_test.py @@ -1,3 +1,4 @@ +import bilby.core.prior import numpy as np import unittest from unittest import mock @@ -312,3 +313,43 @@ def test_log_likelihood_value(self): expected = -6 + np.log(9) actual = self.likelihood.log_likelihood() self.assertEqual(expected, actual) + +class MaximumLikelihoodTest(unittest.TestCase): + def setUp(self): + self.x = np.linspace(0, 10, 30) + + def func(x, m, c, **kwargs): + return m*x + c + + self.m = 2 + self.c = 3 + ytrue = func(self.x, self.m, self.c) + self.sigma = 1 + noise = np.random.normal(0, self.sigma, len(self.x)) + self.yobs = ytrue + noise + self.function = func + self.kwargs = dict(kwarg_1='test_kwarg') + + priors = bilby.core.prior.PriorDict() + priors['m'] = bilby.core.prior.Uniform(minimum=0, maximum=10, name='m') + priors['c'] = bilby.core.prior.Uniform(minimum=0, maximum=10, name='c') + fid = priors.sample() + + self.likelihood = likelihoods.GaussianLikelihood( + x=self.x, y=self.yobs, sigma=self.sigma, function=self.function, kwargs=self.kwargs, priors=priors, + fiducial_parameters=fid) + + def tearDown(self): + del self.x + del self.yobs + del self.sigma + del self.function + del self.kwargs + del self.m + del self.c + del self.likelihood + + def test_maximum_likelihood(self): + maxl_parameters = self.likelihood.find_maximum_likelihood_parameters() + self.assertAlmostEqual(maxl_parameters['m'], self.m, places=0) + self.assertAlmostEqual(maxl_parameters['c'], self.c, places=0) \ No newline at end of file From c824fb155f76d5a9ad99c51bc472e3480a0eb7fb Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 20:57:24 +0100 Subject: [PATCH 11/27] Change metadata logger warning --- redback/transient/transient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redback/transient/transient.py b/redback/transient/transient.py index a92bb745..cac53034 100644 --- a/redback/transient/transient.py +++ b/redback/transient/transient.py @@ -906,7 +906,7 @@ def _set_data(self) -> None: 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 From 22c5463041a73bf786697de120ceaf7c49d67f74 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 21:05:15 +0100 Subject: [PATCH 12/27] Message wording --- redback/transient/transient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redback/transient/transient.py b/redback/transient/transient.py index cac53034..47803382 100644 --- a/redback/transient/transient.py +++ b/redback/transient/transient.py @@ -906,7 +906,7 @@ def _set_data(self) -> None: 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. This is not an error, but a warning. That no metadata could be found online.") + 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 From e738123d5a09d140d9bc8b565af17b8ec259941d Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 23:29:10 +0100 Subject: [PATCH 13/27] Change logging threshold and init to have version info --- redback/__init__.py | 4 ++++ redback/likelihoods.py | 28 ++++++++++++++++++---------- 2 files changed, 22 insertions(+), 10 deletions(-) diff --git a/redback/__init__.py b/redback/__init__.py index ec55a32b..7b24c1df 100644 --- a/redback/__init__.py +++ b/redback/__init__.py @@ -2,3 +2,7 @@ transient_models, utils, photosphere, sed, interaction_processes, constraints, plotting, model_library, simulate_transients from redback.transient import afterglow, kilonova, prompt, supernova, tde from redback.sampler import fit_model +from redback.utils import setup_logger + +__version__ = "1.0.1" +setup_logger(log_level='info') \ No newline at end of file diff --git a/redback/likelihoods.py b/redback/likelihoods.py index cd826515..7d923426 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -6,7 +6,6 @@ from redback.utils import logger from bilby.core.prior import DeltaFunction, Constraint - class _RedbackLikelihood(bilby.Likelihood): def __init__(self, x: np.ndarray, y: np.ndarray, function: callable, kwargs: dict = None, priors=None, @@ -86,8 +85,18 @@ def lnlike_scipy_maximize(self, parameter_list): self.parameters.update(self.get_parameter_dictionary_from_list(parameter_list)) return -self.log_likelihood() - def find_maximum_likelihood_parameters(self, iterations=5, maximization_kwargs=None): - from scipy.optimize import differential_evolution + def find_maximum_likelihood_parameters(self, iterations=5, maximization_kwargs=None, method='Nelder-Mead', + break_threshold=1e-3): + """ + Estimate the maximum likelihood + + :param iterations: Iterations to run the minimizer for before stopping. Default is 5. + :param maximization_kwargs: Any extra keyword arguments passed to the scipy minimize function + :param method: Minimize method to use. Default is 'Nelder-Mead' + :param break_threshold: The threshold for the difference in log likelihood to break the loop. Default is 1e-3. + :return: Dictionary of maximum likelihood parameters + """ + from scipy.optimize import minimize parameter_bounds = self.get_bounds_from_priors(self.priors) if self.priors is None: raise ValueError("Priors must be provided to use this functionality") @@ -99,20 +108,19 @@ def find_maximum_likelihood_parameters(self, iterations=5, maximization_kwargs=N old_fiducial_ln_likelihood = self.log_likelihood() for it in range(iterations): logger.info(f"Optimizing fiducial parameters. Iteration : {it + 1}") - print(f"Optimizing fiducial parameters. Iteration : {it + 1}") - output = differential_evolution( + output = minimize( self.lnlike_scipy_maximize, - bounds=parameter_bounds, x0=updated_parameters_list, - **maximization_kwargs, - ) + bounds=parameter_bounds, + method=method, + **maximization_kwargs,) updated_parameters_list = output['x'] updated_parameters = self.get_parameter_dictionary_from_list(updated_parameters_list) self.parameters.update(updated_parameters) new_fiducial_ln_likelihood = self.log_likelihood_ratio() logger.info(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") - print(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") - if new_fiducial_ln_likelihood - old_fiducial_ln_likelihood < 0.1: + logger.info(f"Updated parameters: {updated_parameters}") + if new_fiducial_ln_likelihood - old_fiducial_ln_likelihood < break_threshold: break old_fiducial_ln_likelihood = new_fiducial_ln_likelihood return updated_parameters From fe56b60a53ef20105c699369f308c8d3e4e5ea85 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 6 Mar 2024 23:54:26 +0100 Subject: [PATCH 14/27] logging and other misc changes --- redback/transient/afterglow.py | 3 ++- redback/transient/transient.py | 2 +- redback/utils.py | 8 +++++--- 3 files changed, 8 insertions(+), 5 deletions(-) diff --git a/redback/transient/afterglow.py b/redback/transient/afterglow.py index 6bd6deb0..4f32ff13 100644 --- a/redback/transient/afterglow.py +++ b/redback/transient/afterglow.py @@ -244,7 +244,8 @@ def _set_data(self) -> None: '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: diff --git a/redback/transient/transient.py b/redback/transient/transient.py index 47803382..c4719dd9 100644 --- a/redback/transient/transient.py +++ b/redback/transient/transient.py @@ -137,7 +137,7 @@ def __init__( 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 diff --git a/redback/utils.py b/redback/utils.py index a2b6908b..b8bb4b33 100644 --- a/redback/utils.py +++ b/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): From e6d5f19e4720d37e20d076caedd71b5526eeac42 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Thu, 7 Mar 2024 00:14:35 +0100 Subject: [PATCH 15/27] Update logger info --- redback/likelihoods.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redback/likelihoods.py b/redback/likelihoods.py index 7d923426..f8bf33ea 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -118,7 +118,7 @@ def find_maximum_likelihood_parameters(self, iterations=5, maximization_kwargs=N updated_parameters = self.get_parameter_dictionary_from_list(updated_parameters_list) self.parameters.update(updated_parameters) new_fiducial_ln_likelihood = self.log_likelihood_ratio() - logger.info(f"Fiducial ln likelihood ratio: {new_fiducial_ln_likelihood:.2f}") + logger.info(f"Current lnlikelihood: {new_fiducial_ln_likelihood:.2f}") logger.info(f"Updated parameters: {updated_parameters}") if new_fiducial_ln_likelihood - old_fiducial_ln_likelihood < break_threshold: break From 10566130f0dbcc6e8e7363c92d8e07b7a10a1782 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Fri, 8 Mar 2024 11:33:08 +0100 Subject: [PATCH 16/27] Some more minor changes --- ...alling_a_model_with_inbuilt_constraints.py | 43 +++++++++++++++++++ redback/transient_models/kilonova_models.py | 2 +- 2 files changed, 44 insertions(+), 1 deletion(-) create mode 100644 examples/calling_a_model_with_inbuilt_constraints.py diff --git a/examples/calling_a_model_with_inbuilt_constraints.py b/examples/calling_a_model_with_inbuilt_constraints.py new file mode 100644 index 00000000..783d02a0 --- /dev/null +++ b/examples/calling_a_model_with_inbuilt_constraints.py @@ -0,0 +1,43 @@ +import bilby.core.prior +from bilby.core.prior import PriorDict, Uniform, Constraint +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd + +import redback +from redback.constraints import csm_constraints + +model = 'csm_interaction' +model_priors = redback.priors.get_priors(model=model) +# model = 'csm_interaction_bolometric' +function = redback.model_library.all_models_dict[model] + +priors = bilby.core.prior.PriorDict(conversion_function=csm_constraints) +priors.update(model_priors) +# priors['photosphere_constraint_1'] = Constraint(0, 1) +# priors['photosphere_constraint_2'] = Constraint(0, 1) +priors['csm_mass'] = 58.0 +priors['mej'] = 46 +priors['vej'] = 5500 +priors['r0'] = 617 +priors['nn'] = 8.8 +priors['delta'] = 0. +priors['rho'] = 19 +priors['eta'] = 2 +priors['redshift'] = 0.16 +samples = pd.DataFrame(priors.sample(10)) +time = np.geomspace(0.01, 500, 500) +redshift = 0.01 + + +for x in range(10): + kwargs = samples.iloc[x] + kwargs['output_format'] = 'magnitude' + kwargs['bands'] = ['lsstg'] + # kwargs['interaction_process'] = None + mag = function(time, **kwargs) + print(mag) + # plt.loglog(time, mag) + plt.plot(time, mag) +plt.gca().invert_yaxis() +plt.show() \ No newline at end of file diff --git a/redback/transient_models/kilonova_models.py b/redback/transient_models/kilonova_models.py index 41fca7c2..3ec3f7e6 100644 --- a/redback/transient_models/kilonova_models.py +++ b/redback/transient_models/kilonova_models.py @@ -1557,7 +1557,7 @@ def two_comp_kne_rosswog_heatingrate(time, redshift, mej_1, vej_1, temperature_f """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame + time_temp = np.geomspace(1e-2, 6e6, 300) # in source frame time_obs = time mej = [mej_1, mej_2] From c3079663e0cb1df9de2f0629f0002a634d5d83aa Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 12:54:36 +0100 Subject: [PATCH 17/27] fractional error likelihood --- redback/likelihoods.py | 65 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 62 insertions(+), 3 deletions(-) diff --git a/redback/likelihoods.py b/redback/likelihoods.py index f8bf33ea..8eefd9e8 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -327,19 +327,78 @@ def log_likelihood(self) -> float: """ return np.nan_to_num(self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma)) +class GaussianLikelihoodWithFractionalNoise(GaussianLikelihood): + def __init__( + self, x: np.ndarray, y: np.ndarray, sigma: Union[float, None, np.ndarray], + function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: + """ + A Gaussian likelihood with noise that is proportional to the model. + The parameters are inferred from the arguments of function + + :type x: np.ndarray + :param y: The y values. + :type y: np.ndarray + :param sigma: The standard deviation of the noise. This is part of the full noise. + The sigma used in the likelihood is sigma = sqrt(sigma_i^2*model_y**2) + :type sigma: Union[float, None, np.ndarray] + :param function: + The python function to fit to the data. Note, this must take the + dependent variable as its first argument. The other arguments + will require a prior and will be sampled over (unless a fixed + value is given). + :type function: callable + :param kwargs: Any additional keywords for 'function'. + :type kwargs: dict + :param priors: The priors for the parameters. Default to None if not provided. + Only necessary if using maximum likelihood estimation functionality. + :type priors: Union[dict, None] + :param fiducial_parameters: The starting guesses for model parameters to + use in the optimization for maximum likelihood estimation. Default to None if not provided. + :type fiducial_parameters: Union[dict, None] + """ + self.sigma = sigma + # These lines of code infer the parameters from the provided function + super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs, priors=priors, + fiducial_parameters=fiducial_parameters) + + @property + def full_sigma(self) -> Union[float, np.ndarray]: + """ + :return: The standard deviation of the full noise + :rtype: Union[float, np.ndarray] + """ + model_y = self.function(self.x, **self.parameters, **self.kwargs) + return np.sqrt(self.sigma_i**2.*model_y**2) + + def noise_log_likelihood(self) -> float: + """ + :return: The noise log-likelihood, i.e. the log-likelihood assuming the signal is just noise. + :rtype: float + """ + if self._noise_log_likelihood is None: + self._noise_log_likelihood = self._gaussian_log_likelihood(res=self.y, sigma=self.sigma) + return self._noise_log_likelihood + + def log_likelihood(self) -> float: + """ + :return: The log-likelihood. + :rtype: float + """ + return np.nan_to_num(self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma)) + class GaussianLikelihoodWithSystematicNoise(GaussianLikelihood): def __init__( self, x: np.ndarray, y: np.ndarray, sigma_i: Union[float, None, np.ndarray], function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """ - A general Gaussian likelihood - the parameters are inferred from the - arguments of function + A Gaussian likelihood with a systematic noise term that is proportional to the model + some additive noise. + The parameters are inferred from the arguments of function :type x: np.ndarray :param y: The y values. :type y: np.ndarray :param sigma_i: The standard deviation of the noise. This is part of the full noise. - The sigma used in the likelihood is sigma = sqrt(sigma_i^2 + sigma^2) + The sigma used in the likelihood is sigma = sqrt(sigma_i^2 + model_y**2*sigma^2) :type sigma_i: Union[float, None, np.ndarray] :param function: The python function to fit to the data. Note, this must take the From 90a026ebbaed8c55fd5522fe60be90d4519479ec Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 13:24:05 +0100 Subject: [PATCH 18/27] Confusion --- redback/likelihoods.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/redback/likelihoods.py b/redback/likelihoods.py index 8eefd9e8..521bd33b 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -335,6 +335,7 @@ def __init__( A Gaussian likelihood with noise that is proportional to the model. The parameters are inferred from the arguments of function + :param x: The x values. :type x: np.ndarray :param y: The y values. :type y: np.ndarray @@ -368,7 +369,7 @@ def full_sigma(self) -> Union[float, np.ndarray]: :rtype: Union[float, np.ndarray] """ model_y = self.function(self.x, **self.parameters, **self.kwargs) - return np.sqrt(self.sigma_i**2.*model_y**2) + return np.sqrt(self.sigma**2.*model_y**2) def noise_log_likelihood(self) -> float: """ @@ -394,6 +395,7 @@ def __init__( A Gaussian likelihood with a systematic noise term that is proportional to the model + some additive noise. The parameters are inferred from the arguments of function + :param x: The x values. :type x: np.ndarray :param y: The y values. :type y: np.ndarray @@ -452,6 +454,7 @@ def __init__( """A general Gaussian likelihood - the parameters are inferred from the arguments of function. Takes into account non-detections with a Uniform likelihood for those points + :param x: The x values. :type x: np.ndarray :param y: The y values. :type y: np.ndarray From 5c1ce3c02437beaddb13ac8eb3e7042cc84a3cb9 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 14:10:46 +0100 Subject: [PATCH 19/27] Add extra functionality to noise types in simulations --- redback/likelihoods.py | 14 +++++++------- redback/simulate_transients.py | 28 ++++++++++++++++++++++++---- 2 files changed, 31 insertions(+), 11 deletions(-) diff --git a/redback/likelihoods.py b/redback/likelihoods.py index 521bd33b..be6bc4eb 100644 --- a/redback/likelihoods.py +++ b/redback/likelihoods.py @@ -329,7 +329,7 @@ def log_likelihood(self) -> float: class GaussianLikelihoodWithFractionalNoise(GaussianLikelihood): def __init__( - self, x: np.ndarray, y: np.ndarray, sigma: Union[float, None, np.ndarray], + self, x: np.ndarray, y: np.ndarray, sigma_i: Union[float, None, np.ndarray], function: callable, kwargs: dict = None, priors=None, fiducial_parameters=None) -> None: """ A Gaussian likelihood with noise that is proportional to the model. @@ -339,9 +339,9 @@ def __init__( :type x: np.ndarray :param y: The y values. :type y: np.ndarray - :param sigma: The standard deviation of the noise. This is part of the full noise. + :param sigma_i: The standard deviation of the noise. This is part of the full noise. The sigma used in the likelihood is sigma = sqrt(sigma_i^2*model_y**2) - :type sigma: Union[float, None, np.ndarray] + :type sigma_i: Union[float, None, np.ndarray] :param function: The python function to fit to the data. Note, this must take the dependent variable as its first argument. The other arguments @@ -357,9 +357,9 @@ def __init__( use in the optimization for maximum likelihood estimation. Default to None if not provided. :type fiducial_parameters: Union[dict, None] """ - self.sigma = sigma + self.sigma_i = sigma_i # These lines of code infer the parameters from the provided function - super().__init__(x=x, y=y, sigma=sigma, function=function, kwargs=kwargs, priors=priors, + super().__init__(x=x, y=y, sigma=sigma_i, function=function, kwargs=kwargs, priors=priors, fiducial_parameters=fiducial_parameters) @property @@ -369,7 +369,7 @@ def full_sigma(self) -> Union[float, np.ndarray]: :rtype: Union[float, np.ndarray] """ model_y = self.function(self.x, **self.parameters, **self.kwargs) - return np.sqrt(self.sigma**2.*model_y**2) + return np.sqrt(self.sigma_i**2.*model_y**2) def noise_log_likelihood(self) -> float: """ @@ -377,7 +377,7 @@ def noise_log_likelihood(self) -> float: :rtype: float """ if self._noise_log_likelihood is None: - self._noise_log_likelihood = self._gaussian_log_likelihood(res=self.y, sigma=self.sigma) + self._noise_log_likelihood = self._gaussian_log_likelihood(res=self.y, sigma=self.sigma_i) return self._noise_log_likelihood def log_likelihood(self) -> float: diff --git a/redback/simulate_transients.py b/redback/simulate_transients.py index fc95e1b5..204389e4 100644 --- a/redback/simulate_transients.py +++ b/redback/simulate_transients.py @@ -14,7 +14,7 @@ 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 @@ def __init__(self, model, parameters, times, model_kwargs, data_points, 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 @@ -81,8 +86,23 @@ def __init__(self, model, parameters, times, model_kwargs, data_points, 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': + output = np.random.normal(true_output, self.noise_term * true_output) + output_error = self.noise_term * true_output + elif noise_type == 'gaussian': + output = np.random.normal(true_output, self.noise_term) + 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: + output = np.random.normal(output, extra_scatter) + output_error = np.sqrt(output_error**2 + extra_scatter**2) + + data['output'] = output + data['output_error'] = output_error self.data = data def save_transient(self, name): From bc67441c9d5a88569c64f6d8ade4c87e9d6eb09f Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 17:14:57 +0100 Subject: [PATCH 20/27] Some more changes --- redback/simulate_transients.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/redback/simulate_transients.py b/redback/simulate_transients.py index 204389e4..68e9214b 100644 --- a/redback/simulate_transients.py +++ b/redback/simulate_transients.py @@ -88,18 +88,21 @@ def __init__(self, model, parameters, times, model_kwargs, data_points, data['true_output'] = true_output if noise_type == 'gaussianmodel': - output = np.random.normal(true_output, self.noise_term * true_output) + 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': - output = np.random.normal(true_output, self.noise_term) + 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: - output = np.random.normal(output, extra_scatter) - output_error = np.sqrt(output_error**2 + extra_scatter**2) + 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 From 16fab7b7677e09f1bd84a78121321717e035b75c Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 17:20:09 +0100 Subject: [PATCH 21/27] Actually use the credible interval level kwarg.... --- redback/plotting.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/redback/plotting.py b/redback/plotting.py index 85666392..60f6a8a3 100644 --- a/redback/plotting.py +++ b/redback/plotting.py @@ -346,7 +346,7 @@ def _plot_lightcurves(self, axes: matplotlib.axes.Axes, times: np.ndarray) -> No axes.plot(times, ys, color=self.random_sample_color, alpha=self.random_sample_alpha, lw=self.linewidth, zorder=self.zorder) elif self.uncertainty_mode == "credible_intervals": - lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=random_ys_list) + lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=random_ys_list, interval=self.credible_interval_level) axes.fill_between( times, lower_bound, upper_bound, alpha=self.uncertainty_band_alpha, color=self.max_likelihood_color) @@ -635,11 +635,11 @@ def plot_lightcurve( elif self.uncertainty_mode == "credible_intervals": if band in self.band_scaling: if self.band_scaling.get("type") == 'x': - lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list) * self.band_scaling.get(band)) + lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list) * self.band_scaling.get(band), interval=self.credible_interval_level) elif self.band_scaling.get("type") == '+': - lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list) + self.band_scaling.get(band)) + lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list) + self.band_scaling.get(band), interval=self.credible_interval_level) else: - lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list)) + lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=np.array(random_ys_list), interval=self.credible_interval_level) axes.fill_between( times - self._reference_mjd_date, lower_bound, upper_bound, alpha=self.uncertainty_band_alpha, color=color_sample) From 357bad5fd70ee411bd1b27990f70957c76a55032 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Sun, 10 Mar 2024 17:21:24 +0100 Subject: [PATCH 22/27] multiband lightcurve also gets custom credible intervals --- redback/plotting.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/redback/plotting.py b/redback/plotting.py index 60f6a8a3..d0b42f55 100644 --- a/redback/plotting.py +++ b/redback/plotting.py @@ -789,7 +789,7 @@ def plot_multiband_lightcurve( axes[ii].plot(times - self._reference_mjd_date, random_ys, color=color_sample, alpha=self.random_sample_alpha, lw=self.linewidth, zorder=self.zorder) elif self.uncertainty_mode == "credible_intervals": - lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=random_ys_list) + lower_bound, upper_bound, _ = redback.utils.calc_credible_intervals(samples=random_ys_list, interval=self.credible_interval_level) axes[ii].fill_between( times - self._reference_mjd_date, lower_bound, upper_bound, alpha=self.uncertainty_band_alpha, color=color_sample) From ee2d2e20bfa1e5636051dc512ee4177949811253 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Mon, 11 Mar 2024 14:18:54 +0100 Subject: [PATCH 23/27] Change maximum time for all kilonova models --- redback/transient_models/kilonova_models.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/redback/transient_models/kilonova_models.py b/redback/transient_models/kilonova_models.py index 3ec3f7e6..6cc62be1 100644 --- a/redback/transient_models/kilonova_models.py +++ b/redback/transient_models/kilonova_models.py @@ -476,7 +476,7 @@ def mosfit_rprocess(time, redshift, mej, vej, kappa, kappa_gamma, temperature_fl cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value dense_resolution = kwargs.get('dense_resolution', 300) - time_temp = np.geomspace(1e-2, 5e6, dense_resolution) # in source frame + time_temp = np.geomspace(1e-2, 7e6, dense_resolution) # in source frame time_obs = time lbols = _mosfit_kilonova_one_component_lbol(time=time_temp, mej=mej, vej=vej) @@ -555,7 +555,7 @@ def mosfit_kilonova(time, redshift, mej_1, vej_1, temperature_floor_1, kappa_1, cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value dense_resolution = kwargs.get('dense_resolution', 300) - time_temp = np.geomspace(1e-2, 5e6, dense_resolution) # in source frame + time_temp = np.geomspace(1e-2, 7e6, dense_resolution) # in source frame time_obs = time mej = [mej_1, mej_2, mej_3] vej = [vej_1, vej_2, vej_3] @@ -979,7 +979,7 @@ def three_component_kilonova_model(time, redshift, mej_1, vej_1, temperature_flo """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame + time_temp = np.geomspace(1e-2, 7e6, 300) # in source frame time_obs = time mej = [mej_1, mej_2, mej_3] @@ -1068,7 +1068,7 @@ def two_component_kilonova_model(time, redshift, mej_1, vej_1, temperature_floor """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-2, 5e6, 300) # in source frame + time_temp = np.geomspace(1e-2, 7e6, 300) # in source frame time_obs = time mej = [mej_1, mej_2] @@ -1381,7 +1381,7 @@ def one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs): """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-3, 5e6, 300) # in source frame + time_temp = np.geomspace(1e-3, 7e6, 300) # in source frame time_obs = time _, temperature, r_photosphere = _one_component_kilonova_model(time_temp, mej, vej, kappa, **kwargs) @@ -1492,7 +1492,7 @@ def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, ye, **kwargs): """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-3, 5e6, 300) # in source frame + time_temp = np.geomspace(1e-3, 7e6, 300) # in source frame time_obs = time _, temperature, r_photosphere = _one_component_kilonova_rosswog_heatingrate(time_temp, mej, vej, ye, **kwargs) @@ -1557,7 +1557,7 @@ def two_comp_kne_rosswog_heatingrate(time, redshift, mej_1, vej_1, temperature_f """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-2, 6e6, 300) # in source frame + time_temp = np.geomspace(1e-2, 7e6, 300) # in source frame time_obs = time mej = [mej_1, mej_2] @@ -1684,7 +1684,7 @@ def metzger_kilonova_model(time, redshift, mej, vej, beta, kappa, **kwargs): """ cosmology = kwargs.get('cosmology', cosmo) dl = cosmology.luminosity_distance(redshift).cgs.value - time_temp = np.geomspace(1e-4, 1e7, 300) # in source frame + time_temp = np.geomspace(1e-4, 7e6, 300) # in source frame time_obs = time bolometric_luminosity, temperature, r_photosphere = _metzger_kilonova_model(time_temp, mej, vej, beta, kappa, **kwargs) From 60f2c787fdc4201113173384caeceb954da30ef4 Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Mon, 11 Mar 2024 14:23:15 +0100 Subject: [PATCH 24/27] Remove nans and zero from spectrum before creating sncosmo source --- redback/sed.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/redback/sed.py b/redback/sed.py index e9fe7187..acedb63a 100644 --- a/redback/sed.py +++ b/redback/sed.py @@ -304,6 +304,11 @@ def get_correct_output_format_from_spectra(time, time_eval, spectra, lambda_arra :param output_format: 'flux', 'magnitude', 'sncosmo_source', 'flux_density' :return: flux, magnitude or SNcosmo TimeSeries Source depending on output format kwarg """ + # 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[10]) + spectra[spectra.value == 0.] = 1e-30 * np.mean(spectra[10]) + source = TimeSeriesSource(phase=time_eval, wave=lambda_array, flux=spectra) if kwargs['output_format'] == 'flux': bands = kwargs['bands'] From 339293be5399efe7d5f5e9c1e37907c377eab4fa Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Mon, 11 Mar 2024 18:00:05 +0100 Subject: [PATCH 25/27] prior latex labels --- .../two_comp_kne_rosswog_heatingrate.prior | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/redback/priors/two_comp_kne_rosswog_heatingrate.prior b/redback/priors/two_comp_kne_rosswog_heatingrate.prior index a1e5f421..c09b3413 100644 --- a/redback/priors/two_comp_kne_rosswog_heatingrate.prior +++ b/redback/priors/two_comp_kne_rosswog_heatingrate.prior @@ -1,9 +1,9 @@ redshift = Uniform(1e-6, 0.1, 'redshift', latex_label = r'$z$') -mej_1 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') -vej_1 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') -ye_1 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}$') -temperature_floor_1 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') -mej_2 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej} }~(M_\odot)$') -vej_2 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}}~(c)$') -ye_2 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}$') -temperature_floor_2 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}}~(\mathrm{K})$') \ No newline at end of file +mej_1 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej}~1}~(M_\odot)$') +vej_1 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}~1}~(c)$') +ye_1 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}~1$') +temperature_floor_1 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}~1}~(\mathrm{K})$') +mej_2 = Uniform(1e-2, 0.05, 'mej', latex_label = r'$M_{\mathrm{ej}~2}~(M_\odot)$') +vej_2 = Uniform(0.05, 0.3, 'vej', latex_label = r'$v_{\mathrm{ej}~2}~(c)$') +ye_2 = Uniform(0.05, 0.4, 'ye', latex_label = r'$Y_{e}~2$') +temperature_floor_2 = LogUniform(100, 6000, 'temperature_floor', latex_label = r'$T_{\mathrm{floor}~2}~(\mathrm{K})$') \ No newline at end of file From 9ba99fdd0f5406ac37d02022107aad45ad118b1e Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Wed, 13 Mar 2024 17:17:32 +0100 Subject: [PATCH 26/27] allow perturbation to heating rate --- redback/transient_models/kilonova_models.py | 23 ++++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/redback/transient_models/kilonova_models.py b/redback/transient_models/kilonova_models.py index 6cc62be1..c7b7ebd5 100644 --- a/redback/transient_models/kilonova_models.py +++ b/redback/transient_models/kilonova_models.py @@ -1422,8 +1422,21 @@ def one_component_kilonova_model(time, redshift, mej, vej, kappa, **kwargs): **kwargs) -def _calc_new_heating_rate(time, mej, electron_fraction, ejecta_velocity): +def _calc_new_heating_rate(time, mej, electron_fraction, ejecta_velocity, **kwargs): + """ + Heating rate prescription following Rosswog and Korobkin 2022 + + :param time: time in seconds + :param mej: ejecta mass in solar masses + :param electron_fraction: electron fraction + :param ejecta_velocity: ejecta velocity in c + :param kwargs: Additional keyword arguments + :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate. + Default is 1.0 i.e., no perturbation. + :return: heating rate in erg/s + """ heating_terms = get_heating_terms(electron_fraction, ejecta_velocity) + heating_rate_perturbation = kwargs.get('heating_rate_perturbation', 1.0) # rescale m0 = mej * solar_mass c1 = np.exp(heating_terms.c1) @@ -1438,7 +1451,7 @@ def _calc_new_heating_rate(time, mej, electron_fraction, ejecta_velocity): term4 = c2 * np.exp(-time/tau2) term5 = c3 * np.exp(-time/tau3) lum_in = term1*term2 + term3 + term4 + term5 - return lum_in*m0 + return lum_in*m0 * heating_rate_perturbation def _one_component_kilonova_rosswog_heatingrate(time, mej, vej, electron_fraction, **kwargs): tdays = time/day_to_s @@ -1454,7 +1467,7 @@ def _one_component_kilonova_rosswog_heatingrate(time, mej, vej, electron_fractio kappa = kappa_from_electron_fraction(electron_fraction) tdiff = np.sqrt(2.0 * kappa * (m0) / (beta * v0 * speed_of_light)) - lum_in = _calc_new_heating_rate(time, mej, electron_fraction, vej) + lum_in = _calc_new_heating_rate(time, mej, electron_fraction, vej, **kwargs) integrand = lum_in * e_th * (time / tdiff) * np.exp(time ** 2 / tdiff ** 2) bolometric_luminosity = np.zeros(len(time)) @@ -1482,6 +1495,8 @@ def one_comp_kne_rosswog_heatingrate(time, redshift, mej, vej, ye, **kwargs): :param kappa: gray opacity :param kwargs: Additional keyword arguments :param temperature_floor: Temperature floor in K (default 4000) + :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate. + Default is 1.0 i.e., no perturbation. :param 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'. @@ -1547,6 +1562,8 @@ def two_comp_kne_rosswog_heatingrate(time, redshift, mej_1, vej_1, temperature_f :param temperature_floor_2: floor temperature of second component :param kappa_2: gray opacity of second component :param kwargs: Additional keyword arguments + :param heating_rate_perturbation: A fudge factor for heating rate to account for uncertainties in the heating rate. + Default is 1.0 i.e., no perturbation. :param 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'. From b7ad7e1ec9c0a18db4c0488739dbf8fb80ad349b Mon Sep 17 00:00:00 2001 From: Nikhil Sarin Date: Thu, 14 Mar 2024 15:26:45 +0100 Subject: [PATCH 27/27] magnitude plotter gets xlim modifiers too --- redback/plotting.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/redback/plotting.py b/redback/plotting.py index d0b42f55..3ab72172 100644 --- a/redback/plotting.py +++ b/redback/plotting.py @@ -391,11 +391,11 @@ class LuminosityPlotter(IntegratedFluxPlotter): class MagnitudePlotter(Plotter): - xlim_low_phase_model_multiplier = 0.9 - xlim_high_phase_model_multiplier = 1.1 - xlim_high_multiplier = 1.2 - ylim_low_magnitude_multiplier = 0.8 - ylim_high_magnitude_multiplier = 1.2 + xlim_low_phase_model_multiplier = KwargsAccessorWithDefault("xlim_low_multiplier", 0.9) + xlim_high_phase_model_multiplier = KwargsAccessorWithDefault("xlim_high_multiplier", 1.1) + xlim_high_multiplier = KwargsAccessorWithDefault("xlim_high_multiplier", 1.2) + ylim_low_magnitude_multiplier = KwargsAccessorWithDefault("ylim_low_multiplier", 0.8) + ylim_high_magnitude_multiplier = KwargsAccessorWithDefault("ylim_high_multiplier", 1.2) ncols = KwargsAccessorWithDefault("ncols", 2) @property