Merge pull request #165 from nikhil-sarin/changes_to_tde_model

Changes to tde model
nikhil-sarin · Jul 31, 2023 · f3d09f3 · f3d09f3
2 parents 5ac72e5 + b9fe887
commit f3d09f3
Show file tree

Hide file tree

Showing 7 changed files with 42 additions and 33 deletions.
diff --git a/docs/acknowledgements.txt b/docs/acknowledgements.txt
@@ -57,14 +57,14 @@ We recommend the following text for citing :code:`redback` in a scientific publi
 
 For using :code:`redback` to fit transients:
 
-- We fit the transient using the open source software package \texttt{Redback}~\cite{sarin_redback},
+- We fit the transient using the open source software package {\sc Redback}~\cite{sarin_redback},
   with the {model_name}~\citep{model reference} model using the {sampler} sampler~\citep{sampler reference} wrapped with bilby~\citep{Ashton+2019}.
 
 For using :code:`redback` to simulate transients:
 
-- We simulate transients using the open source software package \texttt{Redback}~\cite{sarin_redback}, with the {model_name}~\citep{model reference} model.
+- We simulate transients using the open source software package {\sc Redback}~\cite{sarin_redback}, with the {model_name}~\citep{model reference} model.
 
 For using :code:`redback` to download transient data:
 
-- We download and process this data using the API provided by the open source software package \texttt{Redback}~\cite{sarin_redback}.
+- We download and process this data using the API provided by the open source software package {\sc Redback}~\cite{sarin_redback}.
 
diff --git a/optional_requirements.txt b/optional_requirements.txt
@@ -1,5 +1,4 @@
 gwemlightcurves
-sncosmo
 nestle
 sherpa
 PyQt5

diff --git a/redback/likelihoods.py b/redback/likelihoods.py
@@ -176,7 +176,7 @@ def log_likelihood(self) -> float:
         :return: The log-likelihood.
         :rtype: float
         """
-        return self.log_likelihood_x() + self.log_likelihood_y()
+        return np.nan_to_num(self.log_likelihood_x() + self.log_likelihood_y())
 
 
 class GaussianLikelihoodQuadratureNoise(GaussianLikelihood):
@@ -228,7 +228,7 @@ def log_likelihood(self) -> float:
         :return: The log-likelihood.
         :rtype: float
         """
-        return self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma)
+        return np.nan_to_num(self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma))
 
 class GaussianLikelihoodWithSystematicNoise(GaussianLikelihood):
     def __init__(
@@ -280,7 +280,7 @@ def log_likelihood(self) -> float:
         :return: The log-likelihood.
         :rtype: float
         """
-        return self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma)
+        return np.nan_to_num(self._gaussian_log_likelihood(res=self.residual, sigma=self.full_sigma))
 
 class GaussianLikelihoodQuadratureNoiseNonDetections(GaussianLikelihoodQuadratureNoise):
     def __init__(
@@ -337,7 +337,7 @@ def log_likelihood(self) -> float:
         :return: The log-likelihood.
         :rtype: float
         """
-        return self.log_likelihood_y() + self.log_likelihood_upper_limit()
+        return np.nan_to_num(self.log_likelihood_y() + self.log_likelihood_upper_limit())
 
 
 class GRBGaussianLikelihood(GaussianLikelihood):

diff --git a/redback/priors/metzger_tde.prior → redback/priors/cooling_envelope.prior b/redback/priors/metzger_tde.prior → redback/priors/cooling_envelope.prior
diff --git a/...ack/priors/gaussianrise_metzger_tde.prior → ...riors/gaussianrise_cooling_envelope.prior b/...ack/priors/gaussianrise_metzger_tde.prior → ...riors/gaussianrise_cooling_envelope.prior
diff --git a/...gaussianrise_metzger_tde_bolometric.prior → ...ianrise_cooling_envelope_bolometric.prior b/...gaussianrise_metzger_tde_bolometric.prior → ...ianrise_cooling_envelope_bolometric.prior
diff --git a/redback/transient_models/tde_models.py b/redback/transient_models/tde_models.py
@@ -28,7 +28,7 @@ def _analytic_fallback(time, l0, t_0):
 def _semianalytical_fallback():
     pass
 
-def _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
+def _cooling_envelope(mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
     """
     :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
@@ -153,11 +153,13 @@ def _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
                                    'SMBH_accretion_rate', 'time_temp', 'nulnu',
                                    'time_since_fb','tfb', 'lnu', 'envelope_radius', 'envelope_mass',
                                    'rtidal', 'rcirc', 'termination_time', 'termination_time_id'])
-    # constraint_1 = np.min(np.where(Rv < Rcirc/2.))
-    constraint_1 = len(time_temp)
-    constraint_2 = np.min(np.where(Me < 0.0))
+    try:
+        constraint_1 = np.min(np.where(Rv < Rcirc/2.))
+        constraint_2 = np.min(np.where(Me < 0.0))
+    except ValueError:
+        constraint_1 = len(time_temp)
+        constraint_2 = len(time_temp)
     constraint = np.min([constraint_1, constraint_2])
-    constraint_1 = np.min(np.where(Rv < Rcirc/2.))
     termination_time_id = np.min([constraint_1, constraint_2])
     nu = 6.0e14
     expon = 1. / (np.exp(cc.planck * nu / (cc.boltzmann_constant * Teff)) - 1.0)
@@ -178,14 +180,17 @@ def _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
     output.SMBH_accretion_rate = MdotBH[:constraint]
     output.time_temp = time_temp[:constraint]
     output.time_since_fb = output.time_temp - output.time_temp[0]
-    output.termination_time = time_temp[termination_time_id] - tfb
+    if constraint == len(time_temp):
+        output.termination_time = time_temp[-1] - tfb
+    else:
+        output.termination_time = time_temp[termination_time_id] - tfb
     output.termination_time_id = termination_time_id
     output.tfb = tfb
     output.nulnu = nuLnu40[:constraint] * 1e40
     return output
 
-@citation_wrapper('https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
-def metzger_tde(time, redshift,  mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
+@citation_wrapper('https://arxiv.org/abs/2307.15121,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
+def cooling_envelope(time, redshift, mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
     """
     This model is only valid for time after circulation. Use the gaussianrise_metzgertde model for the full lightcurve
 
@@ -205,7 +210,7 @@ def metzger_tde(time, redshift,  mbh_6, stellar_mass, eta, alpha, beta, **kwargs
     :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'
     """
-    output = _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
+    output = _cooling_envelope(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
     cosmology = kwargs.get('cosmology', cosmo)
     dl = cosmology.luminosity_distance(redshift).cgs.value
     time_obs = time
@@ -250,8 +255,8 @@ def metzger_tde(time, redshift,  mbh_6, stellar_mass, eta, alpha, beta, **kwargs
                                                               spectra=spectra, lambda_array=lambda_observer_frame,
                                                               **kwargs)
 
-@citation_wrapper('Sarin and Metzger in prep,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
-def gaussianrise_metzger_tde_bolometric(time, peak_time, sigma_t, mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
+@citation_wrapper('https://arxiv.org/abs/2307.15121,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
+def gaussianrise_cooling_envelope_bolometric(time, peak_time, sigma_t, mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
     """
     Full lightcurve, with gaussian rise till fallback time and then the metzger tde model,
     bolometric version for fitting the bolometric lightcurve
@@ -265,7 +270,7 @@ def gaussianrise_metzger_tde_bolometric(time, peak_time, sigma_t, mbh_6, stellar
     :param kwargs: Additional parameters
     :return luminosity in ergs/s
     """
-    output = _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
+    output = _cooling_envelope(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
     kwargs['binding_energy_const'] = kwargs.get('binding_energy_const', 0.8)
     tfb_sf = calc_tfb(kwargs['binding_energy_const'], mbh_6, stellar_mass)  # source frame
     f1 = pm.gaussian_rise(time=tfb_sf, a_1=1, peak_time=peak_time * cc.day_to_s, sigma_t=sigma_t * cc.day_to_s)
@@ -285,8 +290,9 @@ def gaussianrise_metzger_tde_bolometric(time, peak_time, sigma_t, mbh_6, stellar
     lbol_func = interp1d(full_time, y=full_lbol, fill_value='extrapolate')
     return lbol_func(time*cc.day_to_s)
 
-@citation_wrapper('Sarin and Metzger in prep,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
-def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_mass, eta, alpha, beta, **kwargs):
+
+@citation_wrapper('https://arxiv.org/abs/2307.15121,https://ui.adsabs.harvard.edu/abs/2022arXiv220707136M/abstract')
+def gaussianrise_cooling_envelope(time, redshift, peak_time, sigma_t, 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
@@ -299,6 +305,8 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
     :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'.
@@ -310,12 +318,14 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
     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
-    output = _metzger_tde(mbh_6, stellar_mass, eta, alpha, beta, **kwargs)
+    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.gaussian_rise(time=tfb_obf, a_1=1, peak_time=peak_time * cc.day_to_s, sigma_t=sigma_t * cc.day_to_s)
+    f1 = pm.gaussian_rise(time=stitching_point, a_1=1., peak_time=peak_time * cc.day_to_s, sigma_t=sigma_t * cc.day_to_s)
 
     if kwargs['output_format'] == 'flux_density':
         frequency = kwargs['frequency']
@@ -336,8 +346,8 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
         # build flux density function for each frequency
         flux_den_interp_func = {}
         for freq in unique_frequency:
-            tt_pre_fb = np.linspace(0, tfb_obf / cc.day_to_s, 200) * cc.day_to_s
-            tt_post_fb = output.time_temp * (1 + redshift)
+            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.gaussian_rise(time=tt_pre_fb, a_1=norm_dict[freq],
                                   peak_time=peak_time * cc.day_to_s, sigma_t=sigma_t * cc.day_to_s)
@@ -361,9 +371,9 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
         unique_bands = np.unique(bands)
         temp_kwargs = kwargs.copy()
         temp_kwargs['bands'] = unique_bands
-        f2 = metzger_tde(time=0., redshift=redshift,
-                            mbh_6=mbh_6, stellar_mass=stellar_mass, eta=eta, alpha=alpha, beta=beta,
-                            **temp_kwargs)
+        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
@@ -377,7 +387,7 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
 
         flux_den_interp_func = {}
         for band in unique_bands:
-            tt_pre_fb = np.linspace(0, tfb_obf / cc.day_to_s, 100) * cc.day_to_s
+            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.gaussian_rise(time=tt_pre_fb, a_1=norm_dict[band],
@@ -386,9 +396,9 @@ def gaussianrise_metzger_tde(time, redshift, peak_time, sigma_t, mbh_6, stellar_
                 f1 = calc_ABmag_from_flux_density(f1).value
             temp_kwargs = kwargs.copy()
             temp_kwargs['bands'] = band
-            f2 = metzger_tde(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)
+            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')