SLSN bug fixes (#671)

* Fix i-band issue + bad galactic extinction computation

* Fixed positive K-correction

* Updated model with fixed K-correction

* add dust_extinction

* Add exception to ruff

* Removed comments

---------

Co-authored-by: Etienne Russeil <etru7215@localhost.localdomain>

Author: erusseil

.github/workflows/run_test_ztf.yml

@@ -44,6 +44,7 @@ jobs:
     - name: Run test suites for ZTF
       run: |
         pip install onnxruntime==1.16.3
+        pip install dust_extinction
         rm -f /tmp/forest_*.onnx
         ./run_tests.sh -s ztf
         curl -s https://codecov.io/bash | bash

.ruff.toml

@@ -44,7 +44,8 @@ ignore = [
     "D400", "D401", "D100", "D102", "D103", "D104", "D415", "D419", "D301", 
     "E731",
     "N812", "N806", "N803", 
-    "PLR0913"
+    "PLR0913",
+    "D420"
 ]
 
 # Allow fix for all enabled rules (when `--fix`) is provided.

fink_science/data/models/superluminous_classifier.joblib

@@ -19,7 +19,6 @@
 curdir = os.path.dirname(os.path.abspath(__file__))
 
 classifier_path = curdir + "/data/models/superluminous_classifier.joblib"
-# Declare i-band with dummy value
 band_wave_aa = {1: 4770.0, 2: 6231.0}
 temperature = "sigmoid"
 bolometric = "bazin"

fink_science/ztf/superluminous/kernel.py

@@ -157,17 +157,13 @@ def superluminous_score(
         probas_total = np.zeros(len(objectId), dtype=float) - 1
         transient_mask = pdf["is_transient"]
         old_enough_mask = pdf["jd"] - pdf["jdstarthist"] >= kern.min_duration
-
-        # FIXME: kernel & model do not handle i-band
-        contains_i_band = pdf["cfid"].apply(lambda bands: 3 in bands)
-        mask_valid = transient_mask & old_enough_mask & ~contains_i_band
+        mask_valid = transient_mask & old_enough_mask
 
         if sum(mask_valid) == 0:
             return pd.Series([-1.0] * len(objectId))
 
         # select only transient alerts
         pdf_valid = pdf[mask_valid].copy().reset_index()
-
         valid_ids = list(pdf_valid["objectId"])
 
         # Use Fink API to get the full light curves history
@@ -218,6 +214,9 @@ def superluminous_score(
         lcs = slsn.compute_flux(lcs)
         lcs = slsn.remove_nan(lcs)
 
+        # keep only g and r band, that were used for training
+        lcs = slsn.remove_bad_bands(lcs)
+
         # Perform feature extraction
         features = slsn.extract_features(lcs)
 
@@ -244,7 +243,11 @@ def superluminous_score(
             upper_M = np.array(
                 SLSN_features.apply(
                     lambda x: slsn.abs_peak(
-                        x["peak_mag"], x["photoz"], x["photozerr"], x["ebv"]
+                        [x["peak_mag_g"], x["peak_mag_r"]],
+                        [kern.band_wave_aa[1], kern.band_wave_aa[2]],
+                        x["photoz"],
+                        x["photozerr"],
+                        x["ebv"],
                     )[2],
                     axis=1,
                 )

fink_science/ztf/superluminous/processor.py

@@ -23,6 +23,7 @@
 from fink_science.tester import spark_unit_tests
 from fink_utils.photometry.conversion import mag2fluxcal_snana
 import astropy.units as u
+from dust_extinction.parameter_averages import F99
 from astropy.cosmology import LambdaCDM
 from astropy.coordinates import SkyCoord
 from dustmaps.sfd import SFDQuery
@@ -78,7 +79,38 @@ def compute_flux(pdf):
     return pdf
 
 
-def abs_peak(app_peak, z, zerr, ebv):
+def compute_milky_way_extinction(ebv, lambda_angstrom, Rv=3.1):
+    """Compute the milky way extinction
+
+    Parameters
+    ----------
+    ebv: float
+        E(B-V) extinction.
+    lambda_angstrom: float
+        Effective wavelength of the telescope filter expressed in Angstrom
+    Rv: float
+        Parameter describing the shape of the extinction curve.
+        Rv = 3.1 is a standard value in many cases.
+
+    Examples
+    --------
+    >>> round(compute_milky_way_extinction(0.5, 6000), 2)
+    1.34
+    """
+    # Filter effective wavelength
+    lambda_eff = lambda_angstrom * u.AA
+
+    # Extinction law
+    ext = F99(Rv=Rv)
+
+    # Total extinction
+    R_lambda = ext(lambda_eff) * Rv
+    A_lambda = R_lambda * ebv
+
+    return A_lambda
+
+
+def abs_peak(app_peak, lambda_angstrom, z, zerr, ebv):
     """Compute the peak absolute magnitude based on redshift, assuming a cosmology
 
     Notes
@@ -87,8 +119,11 @@ def abs_peak(app_peak, z, zerr, ebv):
 
     Parameters
     ----------
-    app_peak: float
-        Apparent peak magnitude.
+    app_peak: list
+        Apparent peak magnitudes in each passband
+    lambda_angstrom: list
+        Effective wavelength associated to the peak apparent magnitudes,
+        expressed in Angstrom.
     z: float
         Redshift
     zerr: float
@@ -98,35 +133,54 @@ def abs_peak(app_peak, z, zerr, ebv):
 
     Examples
     --------
-    >>> abs_peak(19, 0.2, 0.05, 0.1)
-    array([-20.49163613, -21.1351084 , -21.63604614])
-    >>> abs_peak(19, 0.2, 0.05, -1)
-    array([-20.18163613, -20.8251084 , -21.32604614])
-    >>> abs_peak(19, 0.2, np.nan, 0.1)
+    >>> abs_peak(19, 4000, 0.2, 0.05, 0.1)
+    array([-20.92638971, -21.66227902, -22.25186059])
+    >>> abs_peak(19, 4000, 0.2, 0.05, -1)
+    array([-20.48512533, -21.22101463, -21.81059621])
+    >>> abs_peak(19, 4000, 0.2, np.nan, 0.1)
     array([ nan,  nan,  nan])
-    >>> abs_peak(19, np.nan, 0.05, 0.1)
+    >>> abs_peak(19, 4000, np.nan, 0.05, 0.1)
     array([ nan,  nan,  nan])
+    >>> abs_peak([18, 18], [4400, 6600], 0.12, 0.01, 0.5)
+    array([-22.74727368, -22.96008329, -23.15747603])
     """
+    # In case the user gives a single value instead of a list
+    app_peak_is_num = (type(app_peak) is float) | (type(app_peak) is int)
+    lambda_angstrom_is_num = (type(lambda_angstrom) is float) | (
+        type(lambda_angstrom) is int
+    )
+
+    if app_peak_is_num & lambda_angstrom_is_num:
+        app_peak = [app_peak]
+        lambda_angstrom = [lambda_angstrom]
+
+    # In case a negative E(B-V) value is provided
     if ebv < 0:
         ebv = 0
 
     if (z == z) and (zerr == zerr):
         cosmo = LambdaCDM(H0=67.8, Om0=0.308, Ode0=0.692)
-        Rv = 3.1
-
-        Ms = []
-        for k in [-1, 0, 1]:
-            effective_z = max(z + k * zerr, 1e-3)
-            D_L = cosmo.luminosity_distance(effective_z).to("pc").value
-            M = (
-                app_peak
-                - 5 * np.log10(D_L / 10)
-                + 2.5 * np.log10(1 + effective_z)
-                - Rv * ebv
-            )
-            Ms.append(M)
 
-        return np.array(Ms)
+        Ms_lambda = []
+
+        for band in range(len(app_peak)):
+            Ms = []
+            for k in [-1, 0, 1]:
+                effective_z = max(z + k * zerr, 1e-3)
+                D_L = cosmo.luminosity_distance(effective_z).to("pc").value
+                M = (
+                    app_peak[band]
+                    - 5 * np.log10(D_L / 10)
+                    - 2.5 * np.log10(1 + effective_z)
+                    - compute_milky_way_extinction(ebv, lambda_angstrom[band])
+                )
+                Ms.append(M)
+            Ms_lambda.append(Ms)
+
+        # Find the band with the highest absolute magnitude
+        brightest = np.argmin(np.array(Ms_lambda)[:, 1])
+
+        return np.array(Ms_lambda[brightest])
 
     return np.array([np.nan, np.nan, np.nan])
 
@@ -271,7 +325,7 @@ def add_all_ebv(pdf):
     Parameters
     ----------
     pdf: pd.DataFrame
-        Must at leat include objectId, ra, dec columns
+        Must at least include objectId, ra, dec columns
 
     Returns
     -------
@@ -331,6 +385,45 @@ def remove_nan(pdf):
     return pdf
 
 
+def remove_bad_bands(pdf):
+    """Keep only the g and r bands
+
+    Parameters
+    ----------
+    pdf: pd.DataFrame
+        Must at least include cfid, based
+        on which it will remove unwanted bands from the columns:
+        "cjd","cmagpsf","csigmapsf","cfid","csigflux","cflux"
+
+    Returns
+    -------
+    pd.DataFrame
+        Original DataFrame with nan/None removed.
+
+    Examples
+    --------
+    >>> pdf = pd.DataFrame(data={"cflux":[[10, 20, 30, 40]],"cfid":[[1, 2, 3, 3]]})
+    >>> result = remove_bad_bands(pdf)
+    >>> expected = pd.DataFrame(data={"cflux":[[10, 20]],"cfid":[[1, 2]]})
+    >>> pd.testing.assert_frame_equal(result, expected)
+    """
+    for k in ["cjd", "cmagpsf", "csigmapsf", "csigflux", "cflux", "cfid"]:
+        if k in pdf.columns:
+            pdf.loc[:, k] = pdf.apply(
+                lambda row: np.array([
+                    a
+                    for a, b in zip(
+                        row[k],
+                        (np.isin(row["cfid"], list(kern.band_wave_aa.keys()))),  # noqa: E711
+                    )
+                    if b
+                ]),
+                axis=1,
+            )
+
+    return pdf
+
+
 def fit_rainbow(lc, rainbow_model):
     """Perform a rainbow fit (Russeil et al. 2024) on a light curve.
 
@@ -364,9 +457,6 @@ def fit_rainbow(lc, rainbow_model):
         np.array(lc["cfid"]),
     )
 
-    # t_scaler = Scaler.from_time(lc["cjd"])
-    # m_scaler = MultiBandScaler.from_flux(lc["cflux"], lc["cfid"], with_baseline=False)
-
     try:
         result, errors = rainbow_model._eval_and_get_errors(
             t=lc["cjd"],
@@ -443,7 +533,7 @@ def statistical_features(lc):
     -------
     list
         List of statistical features
-        [amplitude, kurtosis, max_slope, skew, peak_magn, std_flux, q15_time, q85_time]
+        [amplitude, kurtosis, max_slope, skew, peak_mag_g, peak_mag_r, std_flux, q15_time, q85_time]
     """
     amplitude = lcpckg.Amplitude()
     kurtosis = lcpckg.Kurtosis()
@@ -464,11 +554,14 @@ def statistical_features(lc):
 
     normed_flux = lc["cflux"] / np.max(lc["cflux"])
     shifted_time = lc["cjd"] - np.min(lc["cjd"])
-    peak_mag = np.min(lc["cmagpsf"])
+
+    peak_mag_g = np.min(lc["cmagpsf"][lc["cfid"] == 1], initial=99)
+    peak_mag_r = np.min(lc["cmagpsf"][lc["cfid"] == 2], initial=99)
+
     std = np.std(normed_flux)
     q15 = np.quantile(shifted_time, 0.15)
     q85 = np.quantile(shifted_time, 0.85)
-    return list(result) + [peak_mag, std, q15, q85]
+    return list(result) + [peak_mag_g, peak_mag_r, std, q15, q85]
 
 
 def quiet_model():
@@ -562,17 +655,18 @@ def extract_features(data):
     >>> salt_features = quiet_fit_salt(lc, salt_model)
 
     # Check their values
-    >>> np.testing.assert_allclose(stat_features,[   8.307904e+02,
+    >>> np.testing.assert_allclose(stat_features,[  8.307904e+02,
     ... 4.843807e-02,   7.573933e+03,  -7.161292e-01,
-    ... 1.875300e+01,   1.383518e-01,   9.992026e+00,   2.499306e+01], rtol=1e-3)
+    ... 1.875300e+01,   1.882850e+01,   1.383518e-01,
+    ... 9.992026e+00,   2.499306e+01], rtol=1e-3)
     >>> np.testing.assert_allclose(salt_features,[  1.374512e-01,
     ... -1.201602e+01,   3.522748e-03,   9.219506e+00,
     ... 3.321469e-02,   4.337947e+01], rtol=5e-2)
     >>> np.testing.assert_allclose(rainbow_features,
-    ... [ -2.161259e+00,   4.886508e+03,   2.196836e+01,   2.740976e+01,
-    ... 9.102432e+03,   9.948595e+03,   1.403806e+00,  -5.663001e-01,
-    ... 1.050990e+01,   6.421245e+00,   1.106539e+00,   7.157673e+00,
-    ... 1.364669e+01,   1.184238e+00,   1.194966e-01], rtol=5e-2)
+    ... [ -2.161261e+00,   4.886507e+03,   2.196836e+01,   2.740982e+01,
+    ... 9.102431e+03,   9.948591e+03,   1.403805e+00,  -5.663014e-01,
+    ... 1.050993e+01,   6.421246e+00,   1.106546e+00,   7.157723e+00,
+    ... 1.364673e+01,   1.184242e+00,   1.194966e-01], rtol=5e-2)
 
     # Check full feature extraction function
     >>> pdf_check = pdf.copy()
@@ -581,15 +675,15 @@ def extract_features(data):
     # Only the fake alert should pass the cuts
     >>> np.testing.assert_equal(
     ... np.array(np.sum(full_features.iloc[-30:].isnull(), axis=1)),
-    ... np.array([ 29, 29, 29,  29, 29, 29,  29, 29, 29,  29, 29, 29,  29,
-    ... 29, 29, 29,  29, 29, 29,  29, 29, 29, 29, 29, 29,  29, 29, 29, 29, 0]))
+    ... np.array([ 30, 30, 30,  30, 30, 30,  30, 30, 30,  30, 30, 30,  30,
+    ... 30, 30, 30,  30, 30, 30,  30, 30, 30, 30, 30, 30,  30, 30, 30, 30, 0]))
 
     >>> list(full_features.columns) == ["distnr", "ra", "dec", "ebv", "duration",
-    ... "flux_amplitude", "kurtosis", "max_slope", "skew", "peak_mag", "std_flux", "q15",
-    ... "q85", "reference_time", "amplitude", "rise_time", "fall_time", "Tmin",
-    ... "Tmax", "t_color", "snr_reference_time", "snr_amplitude", "snr_rise_time",
-    ... "snr_fall_time", "snr_Tmin", "snr_Tmax", "snr_t_color", "chi2_rainbow",
-    ... "z", "t0", "x0", "x1", "c", "chi2_salt"]
+    ... "flux_amplitude", "kurtosis", "max_slope", "skew", "peak_mag_g", "peak_mag_r",
+    ... "std_flux", "q15", "q85", "reference_time", "amplitude", "rise_time", "fall_time",
+    ... "Tmin", "Tmax", "t_color", "snr_reference_time", "snr_amplitude", "snr_rise_time",
+    ... "snr_fall_time", "snr_Tmin", "snr_Tmax", "snr_t_color", "chi2_rainbow", "z", "t0",
+    ... "x0", "x1", "c", "chi2_salt"]
     True
     """
     data = add_all_ebv(data)
@@ -617,7 +711,8 @@ def extract_features(data):
             "kurtosis",
             "max_slope",
             "skew",
-            "peak_mag",
+            "peak_mag_g",
+            "peak_mag_r",
             "std_flux",
             "q15",
             "q85",
@@ -636,6 +731,7 @@ def extract_features(data):
             kern.min_points_perband
             <= np.array([sum(lc["cfid"] == band) for band in list(kern.band_wave_aa)])
         )
+
         enough_total_points = len(lc["cjd"]) > kern.min_points_total
         duration = np.ptp(lc["cjd"])
         enough_duration = duration > kern.min_duration

fink_science/ztf/superluminous/slsn_classifier.py

@@ -44,6 +44,7 @@ light-curve==0.10.4
 sncosmo==2.12.1
 xgboost==2.1.4
 dustmaps==1.0.14
+dust_extinction==1.5
 
 #for ELAsTiCC
 light-curve[full]