Add kilonova event source type

requirements.txt

@@ -10,6 +10,8 @@ colossus
 speclite
 pyyaml
 matplotlib
+redback>=1.17.0
+colossus
 uncertainties
 kcorrect
 tqdm

slsim/Sources/Events/BNSMerger/kilonova.py

@@ -0,0 +1,133 @@
+import numpy as np
+
+if not hasattr(np, "trapezoid"):
+    np.trapezoid = np.trapz
+
+from astropy import cosmology
+from redback.transient_models import kilonova_models
+
+
+class Kilonova:
+    """Class for initializing a kilonova light curve model.
+
+    If modeldir is provided, external kilonova model files are expected.
+    This option is currently not supported. If modeldir is not provided,
+    the model is retrieved from Redback's built-in kilonova models. By
+    default, the MOSFiT-based kilonova model is used. Information about
+    Redback can be found at
+    https://redback.readthedocs.io/en/latest/.
+    """
+
+    def __init__(
+        self,
+        redshift,
+        model_name="mosfit_kilonova",
+        ejecta_mass=None,
+        ejecta_velocity=None,
+        opacity=None,
+        temperature_floor=None,
+        kappa_gamma=10,
+        mag_zpsys="AB",
+        cosmo=cosmology.FlatLambdaCDM(H0=70, Om0=0.3),
+        modeldir=None,
+        **kwargs,
+    ):
+        """
+        :param redshift: The redshift of the kilonova source.
+        :type redshift: float
+        :param model_name: The kilonova light curve model to be used. If not provided,
+            the default model is the MOSFiT-based kilonova model.
+        :type model_name: str
+        :param ejecta_mass: Ejecta masses for the kilonova components.
+        :type ejecta_mass: array-like or None
+        :param ejecta_velocity: Ejecta velocities for the kilonova components.
+        :type ejecta_velocity: array-like or None
+        :param opacity: Opacities for the kilonova components.
+        :type opacity: array-like or None
+        :param temperature_floor: Temperature floors for the kilonova components.
+        :type temperature_floor: array-like or None
+        :param kappa_gamma: Gamma-ray opacity.
+        :type kappa_gamma: float
+        :param mag_zpsys: Optional, AB or Vega (AB default).
+        :type mag_zpsys: str
+        :param cosmo: Cosmology for luminosity distance calculation.
+        :type cosmo: `~astropy.cosmology`
+        :param modeldir: Directory including files for external kilonova models.
+        :type modeldir: str or None
+        :param kwargs: Additional keyword arguments passed to the Redback kilonova model.
+        :type kwargs: dict
+        """
+
+        if modeldir is not None:
+            # external kilonova model
+            raise NotImplementedError(
+                "External kilonova model files are not supported yet."
+            )
+        else:
+            # use Redback built-in kilonova model, e.g. mosfit_kilonova
+            if not hasattr(kilonova_models, model_name):
+                raise ValueError(
+                    f"{model_name} is not available in "
+                    "redback.transient_models.kilonova_models."
+                )
+            else:
+                self._model = getattr(kilonova_models, model_name)
+
+        self._model_name = model_name
+        self._redshift = redshift
+        self._mag_zpsys = mag_zpsys
+        self._cosmo = cosmo
+        self._kwargs = kwargs
+
+        parameter_groups = {
+            "ejecta_mass": ejecta_mass,
+            "ejecta_velocity": ejecta_velocity,
+            "opacity": opacity,
+            "temperature_floor": temperature_floor,
+        }
+
+        for name, values in parameter_groups.items():
+            if values is None:
+                raise ValueError(f"{name} must be provided.")
+            if len(values) != 3:
+                raise ValueError(f"{name} must have three components.")
+
+        self._model_parameters = {
+            "mej_1": ejecta_mass[0],
+            "mej_2": ejecta_mass[1],
+            "mej_3": ejecta_mass[2],
+            "vej_1": ejecta_velocity[0],
+            "vej_2": ejecta_velocity[1],
+            "vej_3": ejecta_velocity[2],
+            "kappa_1": opacity[0],
+            "kappa_2": opacity[1],
+            "kappa_3": opacity[2],
+            "temperature_floor_1": temperature_floor[0],
+            "temperature_floor_2": temperature_floor[1],
+            "temperature_floor_3": temperature_floor[2],
+            "kappa_gamma": kappa_gamma,
+        }
+
+    def get_apparent_magnitude(self, time, band, zpsys="AB"):
+        """Function to return apparent magnitude of a kilonova for a given band
+        and time.
+
+        :param time: The observer-frame time array to evaluate the model
+            (in days)
+        :type time: array-like
+        :param band: The band to evaluate the model over.
+        :type band: str or list
+        :param zpsys: Optional, AB or Vega (AB default)
+        :type zpsys: str
+        :return: magnitude of source
+        """
+
+        return self._model(
+            time=time,
+            redshift=self._redshift,
+            bands=band,
+            output_format="magnitude",
+            cosmology=self._cosmo,
+            **self._model_parameters,
+            **self._kwargs,
+        )

slsim/Sources/SourceTypes/kilonova_event.py

@@ -0,0 +1,145 @@
+import warnings
+from slsim.Sources.SourceTypes.source_base import SourceBase
+from slsim.Sources.Events.BNSMerger.kilonova import Kilonova
+from slsim.ImageSimulation.image_quality_lenstronomy import (
+    get_all_supported_bands,
+    get_sncosmo_filtername,
+)
+
+
+class KilonovaEvent(SourceBase):
+    """A class to manage a BNS merger."""
+
+    def __init__(
+        self,
+        lightcurve_time,
+        variability_model,
+        model_name="mosfit_kilonova",
+        mag_zpsys="AB",
+        modeldir=None,
+        kwargs_variability=None,
+        kwargs_kilonova=None,
+        cosmo=None,
+        **kwargs,
+    ):
+        """
+        :param lightcurve_time: Observation time array for the light curve in units of days.
+        :type lightcurve_time: array-like
+        :param variability_model: Keyword for the variability model to be used. This is an
+            input for the Variability class.
+        :type variability_model: str
+        :param model_name: Kilonova light curve model to be used. If not provided, the
+            default model is the MOSFiT-based kilonova model.
+        :type model_name: str
+        :param mag_zpsys: Optional, AB or Vega (AB default).
+        :type mag_zpsys: str
+        :param modeldir: Directory including files for external kilonova models. This
+            option is currently not supported.
+        :type modeldir: str or None
+        :param kwargs_variability: Dictionary with bands as strings, each containing
+            input configurations for point source variability.
+        :type kwargs_variability: dict of dict or None
+        :param kwargs_kilonova: Keyword arguments passed to the Kilonova class, such as
+            ejecta_mass, ejecta_velocity, opacity, temperature_floor, and kappa_gamma.
+        :type kwargs_kilonova: dict or None
+        :param cosmo: Astropy cosmology instance.
+        :type cosmo: `~astropy.cosmology`
+        :param kwargs: Keyword arguments passed to the SourceBase class. This may contain
+            source properties such as redshift and offsets from the host galaxy.
+        :type kwargs: dict
+        """
+        super().__init__(
+            extended_source=False,
+            point_source=True,
+            cosmo=cosmo,
+            variability_model=variability_model,
+            **kwargs,
+        )
+        self.name = "BNS"
+        self._variability_computed = False
+        self._kwargs_variability = kwargs_variability
+        self._lightcurve_time = lightcurve_time
+
+        self._model_name = model_name
+        self._mag_zpsys = mag_zpsys
+        self._modeldir = modeldir
+        self._kwargs_kilonova = kwargs_kilonova or {}
+
+    @property
+    def light_curve(self):
+        """Provides lightcurves of a bns merger in each band."""
+        if self._kwargs_variability is not None:
+            kwargs_variab_extracted = {}
+            if self._cosmo is None:
+                raise ValueError(
+                    "Cosmology cannot be None for BNSMerger class. Please"
+                    "provide a suitable astropy cosmology."
+                )
+            else:
+                # Initialize BNS/Kilonova light curve model
+                lightcurve_class = Kilonova(
+                    redshift=self._z,
+                    model_name=self._model_name,
+                    mag_zpsys=self._mag_zpsys,
+                    cosmo=self._cosmo,
+                    modeldir=self._modeldir,
+                    **self._kwargs_kilonova,
+                )
+                self._lightcurve_class = lightcurve_class
+
+            supported_bands = get_all_supported_bands()
+            provided_bands = set(supported_bands) & set(self._kwargs_variability)
+
+            for element in provided_bands:
+                name = "ps_mag_" + element
+                times = self._lightcurve_time
+
+                # Use the sncosmo band-name mapping since Redback expects registered filter names.
+                provided_band = get_sncosmo_filtername(element)
+
+                try:
+                    magnitudes = lightcurve_class.get_apparent_magnitude(
+                        time=times,
+                        band=provided_band,
+                        zpsys=self._mag_zpsys,
+                    )
+                except Exception as e:
+                    warnings.warn(
+                        f"Skipping band '{provided_band}': Failed to generate lightcurve. "
+                        f"(Error: {e})",
+                        UserWarning,
+                    )
+                    continue
+
+                if name not in self.source_dict:
+                    self.source_dict[name] = float(min(magnitudes))
+
+                kwargs_variab_extracted[element] = {
+                    "MJD": times,
+                    name: magnitudes,
+                }
+        else:
+            kwargs_variab_extracted = {}
+
+        self._variability_computed = True
+        return kwargs_variab_extracted
+
+    def point_source_magnitude(self, band, image_observation_times=None):
+        """Get the magnitude of the BNS/kilonova point source in a specific
+        band.
+
+        :param band: Imaging band.
+        :type band: str
+        :param image_observation_times: Image observation times. If
+            None, takes the peak magnitude.
+        :type image_observation_times: array-like or None
+        :return: Magnitude of the point source in the specified band.
+        :rtype: float or array-like
+        """
+
+        if not self._variability_computed:
+            self._kwargs_variability_model = self.light_curve
+
+        return super().point_source_magnitude(
+            band=band, image_observation_times=image_observation_times
+        )

slsim/Sources/source.py

@@ -2,7 +2,7 @@
 from slsim.Sources.SourceTypes.source_base import SourceBase
 from copy import deepcopy
 
-_SUPPORTED_POINT_SOURCES = ["supernova", "quasar", "general_lightcurve"]
+_SUPPORTED_POINT_SOURCES = ["supernova", "quasar", "general_lightcurve", "kilonova"]
 _SUPPORTED_EXTENDED_SOURCES = [
     "single_sersic",
     "double_sersic",
@@ -27,7 +27,7 @@ def __init__(
          extended source types are 'single_sersic', 'double_sersic', 'catalog_source', and 'interpolated'.
         :type extended_source_type: str or None
         :param point_source_type: Keyword to specify type of point source. Supported point
-         source types are 'supernova', 'quasar', and 'general_lightcurve'.
+         source types are 'supernova', 'quasar', 'general_lightcurve', and 'kilonova'.
         :type point_source_type: str or None
         :param source_dict: Source properties. Can be a dictionary or an Astropy table.
          For a detailed description of this dictionary, please see the documentation for
@@ -69,6 +69,10 @@ def __init__(
             from slsim.Sources.SourceTypes.general_lightcurve import GeneralLightCurve
 
             self._source = GeneralLightCurve(**source_dict)
+        elif source_type in ["kilonova"]:
+            from slsim.Sources.SourceTypes.kilonova_event import KilonovaEvent
+
+            self._source = KilonovaEvent(**source_dict)
 
         # extended sources
         elif source_type in ["single_sersic"]:

tests/test_Sources/test_Events/test_BNSMerger/test_kilonova.py

@@ -0,0 +1,78 @@
+import numpy as np
+from slsim.Sources.Events.BNSMerger.kilonova import Kilonova
+import numpy.testing as npt
+import pytest
+
+
+@pytest.fixture
+def Kilonova_class():
+    KN = Kilonova(
+        redshift=0.1,
+        model_name="mosfit_kilonova",
+        ejecta_mass=[0.01, 0.02, 0.03],
+        ejecta_velocity=[0.1, 0.2, 0.3],
+        opacity=[0.5, 3.0, 10.0],
+        temperature_floor=[5000, 4000, 3000],
+        mag_zpsys="AB",
+        dense_resolution=50,
+    )
+
+    return KN
+
+
+def test_kilonova_mag(Kilonova_class):
+    time = np.array([0.5, 1.0, 2.0])
+    mag = Kilonova_class.get_apparent_magnitude(time=time, band="lsstr")
+
+    npt.assert_equal(np.shape(mag), np.shape(time))
+    npt.assert_(np.all(np.isfinite(mag)))
+    npt.assert_(np.all(mag > 0))
+
+
+def test_kilonova_missing_parameters():
+    with pytest.raises(ValueError):
+        Kilonova(
+            redshift=0.1,
+            ejecta_velocity=[0.1, 0.2, 0.3],
+            opacity=[0.5, 3.0, 10.0],
+            temperature_floor=[5000, 4000, 3000],
+        )
+
+
+def test_kilonova_parameter_length():
+    with pytest.raises(ValueError):
+        Kilonova(
+            redshift=0.1,
+            ejecta_mass=[0.01, 0.02],
+            ejecta_velocity=[0.1, 0.2, 0.3],
+            opacity=[0.5, 3.0, 10.0],
+            temperature_floor=[5000, 4000, 3000],
+        )
+
+
+def test_kilonova_invalid_model_name():
+    with pytest.raises(ValueError):
+        Kilonova(
+            redshift=0.1,
+            model_name="not_a_kilonova_model",
+            ejecta_mass=[0.01, 0.02, 0.03],
+            ejecta_velocity=[0.1, 0.2, 0.3],
+            opacity=[0.5, 3.0, 10.0],
+            temperature_floor=[5000, 4000, 3000],
+        )
+
+
+def test_kilonova_external_modeldir_not_supported():
+    with pytest.raises(NotImplementedError):
+        Kilonova(
+            redshift=0.1,
+            ejecta_mass=[0.01, 0.02, 0.03],
+            ejecta_velocity=[0.1, 0.2, 0.3],
+            opacity=[0.5, 3.0, 10.0],
+            temperature_floor=[5000, 4000, 3000],
+            modeldir="some/path",
+        )
+
+
+if __name__ == "__main__":
+    pytest.main()