Improved initial guesses/limits of Rainbow functions (#494)

* Improved initial guesses/limits of Rainbow functions

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Changed the tests from parameter comparison to flux comparison

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Reformat file

* reformat again :)

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Etienne Russeil <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

light-curve/light_curve/light_curve_py/features/rainbow/bolometric.py

@@ -102,12 +102,12 @@ def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
         m_amplitude = np.ptp(m)
 
-        mean_dt = np.median(t[1:] - t[:-1])
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         limits = {}
         limits["reference_time"] = (np.min(t) - 10 * t_amplitude, np.max(t) + 10 * t_amplitude)
         limits["amplitude"] = (0.0, 20 * m_amplitude)
-        limits["rise_time"] = (0.1 * mean_dt, 10 * t_amplitude)
+        limits["rise_time"] = (dt / 100, 10 * t_amplitude)
 
         return limits
 
@@ -147,24 +147,13 @@ def value(t, t0, amplitude, rise_time, fall_time):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        A = np.ptp(m)
-
-        mc = m - np.min(m)  # To avoid crashing on all-negative data
+        A = 1.5 * max(np.max(m), np.ptp(m))
 
-        # Naive peak position from the highest point
-        t0 = t[np.argmax(m)]
-        # Peak position as weighted centroid of everything above median
-        idx = m > np.median(m)
-        # t0 = np.sum(t[idx] * m[idx] / sigma[idx]) / np.sum(m[idx] / sigma[idx])
-        # Weighted centroid sigma
-        dt = np.sqrt(np.sum((t[idx] - t0) ** 2 * (mc[idx]) / sigma[idx]) / np.sum(mc[idx] / sigma[idx]))
+        t0, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         # Empirical conversion of sigma to rise/fall times
-        rise_time = dt / 2
-        fall_time = dt / 2
-
-        # Compensate for the difference between reference_time and peak position
-        t0 -= np.log(fall_time / rise_time) * rise_time * fall_time / (rise_time + fall_time)
+        rise_time = dt
+        fall_time = dt
 
         initial = {}
         initial["reference_time"] = t0
@@ -178,14 +167,13 @@ def initial_guesses(t, m, sigma, band):
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
         m_amplitude = np.ptp(m)
-
-        mean_dt = np.median(t[1:] - t[:-1])
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         limits = {}
         limits["reference_time"] = (np.min(t) - 10 * t_amplitude, np.max(t) + 10 * t_amplitude)
         limits["amplitude"] = (0.0, 20 * m_amplitude)
-        limits["rise_time"] = (0.1 * mean_dt, 10 * t_amplitude)
-        limits["fall_time"] = (0.1 * mean_dt, 10 * t_amplitude)
+        limits["rise_time"] = (dt / 100, 10 * t_amplitude)
+        limits["fall_time"] = (dt / 100, 10 * t_amplitude)
 
         return limits
 
@@ -198,7 +186,7 @@ def peak_time(t0, amplitude, rise_time, fall_time):
 class LinexpBolometricTerm(BaseBolometricTerm):
     """Linexp function, symmetric form. Generated using a prototype version of Multi-view
     Symbolic Regression (Russeil et al. 2024, https://arxiv.org/abs/2402.04298) on
-    a SLSN ZTF light curve (https://ztf.snad.space/dr17/view/821207100004043)"""
+    a SLSN ZTF light curve (https://ztf.snad.space/dr17/view/821207100004043). Careful not very stable guesses/limits"""
 
     @staticmethod
     def parameter_names():
@@ -226,6 +214,7 @@ def value(t, t0, amplitude, rise_time):
     def initial_guesses(t, m, sigma, band):
         A = np.ptp(m)
         med_dt = median_dt(t, band)
+        t0, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         # Compute points after or before maximum
         peak_time = t[np.argmax(m)]
@@ -276,7 +265,6 @@ def parameter_scalings():
     @staticmethod
     def value(t, t0, amplitude, time1, time2, p):
         dt = t - t0
-
         result = np.zeros_like(dt)
 
         # To avoid numerical overflows
@@ -290,37 +278,34 @@ def value(t, t0, amplitude, time1, time2, p):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        A = np.ptp(m)
-        med_dt = median_dt(t, band)
+        A = max(np.max(m), np.ptp(m))
+        t0, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
-        # Naive peak position from the highest point
-        t0 = t[np.argmax(m)]
-
-        # Empirical conversion of sigma to rise/fall times
-        time1 = 50 * med_dt
-        time2 = 50 * med_dt
+        # Empirical conversion of sigma to times
+        time1 = 2 * dt
+        time2 = 2 * dt
 
         initial = {}
         initial["reference_time"] = t0
         initial["amplitude"] = A
         initial["time1"] = time1
         initial["time2"] = time2
-        initial["p"] = 0.1
+        initial["p"] = 1
 
         return initial
 
     @staticmethod
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
         m_amplitude = np.ptp(m)
-        med_dt = median_dt(t, band)
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         limits = {}
         limits["reference_time"] = (np.min(t) - 10 * t_amplitude, np.max(t) + 10 * t_amplitude)
         limits["amplitude"] = (0.0, 10 * m_amplitude)
-        limits["time1"] = (med_dt, 2 * t_amplitude)
-        limits["time2"] = (med_dt, 2 * t_amplitude)
-        limits["p"] = (1e-4, 10)
+        limits["time1"] = (dt / 10, 2 * t_amplitude)
+        limits["time2"] = (dt / 10, 2 * t_amplitude)
+        limits["p"] = (1e-2, 100)
 
         return limits
 
@@ -336,13 +321,27 @@ def peak_time(t0, p):
 
 def median_dt(t, band):
     # Compute the median distance between points in each band
+    # Caution when using this method as it might be strongly biaised because of ZTF high cadence a given day.
     dt = []
     for b in np.unique(band):
         dt += list(t[band == b][1:] - t[band == b][:-1])
     med_dt = np.median(dt)
     return med_dt
 
 
+def t0_and_weighted_centroid_sigma(t, m, sigma):
+    # To avoid crashing on all-negative data
+    mc = m - np.min(m)
+
+    # Peak position as weighted centroid of everything above median
+    idx = m > np.median(m)
+    t0 = np.sum(t[idx] * m[idx] / sigma[idx]) / np.sum(m[idx] / sigma[idx])
+
+    # Weighted centroid sigma
+    dt = np.sqrt(np.sum((t[idx] - t0) ** 2 * (mc[idx]) / sigma[idx]) / np.sum(mc[idx] / sigma[idx]))
+    return t0, dt
+
+
 bolometric_terms = {
     "sigmoid": SigmoidBolometricTerm,
     "bazin": BazinBolometricTerm,

light-curve/light_curve/light_curve_py/features/rainbow/temperature.py

@@ -76,7 +76,7 @@ def initial_guesses(t, m, sigma, band):
     @staticmethod
     def limits(t, m, sigma, band):
         limits = {}
-        limits["T"] = (1e2, 2e6)  # K
+        limits["T"] = (1e3, 2e6)  # K
 
         return limits
 
@@ -111,24 +111,24 @@ def value(t, t0, temp_min, temp_max, t_color):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        med_dt = median_dt(t, band)
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         initial = {}
         initial["Tmin"] = 7000.0
         initial["Tmax"] = 10000.0
-        initial["t_color"] = 10 * med_dt
+        initial["t_color"] = 2 * dt
 
         return initial
 
     @staticmethod
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
-        med_dt = median_dt(t, band)
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         limits = {}
         limits["Tmin"] = (1e3, 2e6)  # K
         limits["Tmax"] = (1e3, 2e6)  # K
-        limits["t_color"] = (2 * med_dt, 10 * t_amplitude)
+        limits["t_color"] = (dt / 3, 10 * t_amplitude)
 
         return limits
 
@@ -163,25 +163,25 @@ def value(t, t0, Tmin, Tmax, t_color, t_delay):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        med_dt = median_dt(t, band)
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         initial = {}
         initial["Tmin"] = 7000.0
         initial["Tmax"] = 10000.0
-        initial["t_color"] = 10 * med_dt
+        initial["t_color"] = 2 * dt
         initial["t_delay"] = 0.0
 
         return initial
 
     @staticmethod
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
-        med_dt = median_dt(t, band)
+        _, dt = t0_and_weighted_centroid_sigma(t, m, sigma)
 
         limits = {}
         limits["Tmin"] = (1e3, 2e6)  # K
         limits["Tmax"] = (1e3, 2e6)  # K
-        limits["t_color"] = (2 * med_dt, 10 * t_amplitude)
+        limits["t_color"] = (dt / 3, 10 * t_amplitude)
         limits["t_delay"] = (-t_amplitude, t_amplitude)
 
         return limits
@@ -196,6 +196,19 @@ def median_dt(t, band):
     return med_dt
 
 
+def t0_and_weighted_centroid_sigma(t, m, sigma):
+    # To avoid crashing on all-negative data
+    mc = m - np.min(m)
+
+    # Peak position as weighted centroid of everything above median
+    idx = m > np.median(m)
+    t0 = np.sum(t[idx] * m[idx] / sigma[idx]) / np.sum(m[idx] / sigma[idx])
+
+    # Weighted centroid sigma
+    dt = np.sqrt(np.sum((t[idx] - t0) ** 2 * (mc[idx]) / sigma[idx]) / np.sum(mc[idx] / sigma[idx]))
+    return t0, dt
+
+
 temperature_terms = {
     "constant": ConstantTemperatureTerm,
     "sigmoid": SigmoidTemperatureTerm,

light-curve/tests/light_curve_py/features/test_rainbow.py

@@ -15,10 +15,10 @@ def test_noisy_with_baseline():
     fall_time = 30.0
     Tmin = 5e3
     Tmax = 15e3
-    k_sig = 4.0
+    t_color = 10
     baselines = {b: 0.3 * amplitude + rng.exponential(scale=0.3 * amplitude) for b in band_wave_aa}
 
-    expected = [reference_time, amplitude, rise_time, fall_time, Tmin, Tmax, k_sig, *baselines.values(), 1.0]
+    expected = [reference_time, amplitude, rise_time, fall_time, Tmin, Tmax, t_color, *baselines.values(), 1.0]
 
     feature = RainbowFit.from_angstrom(band_wave_aa, with_baseline=True, temperature="sigmoid", bolometric="bazin")
 
@@ -40,7 +40,7 @@ def test_noisy_with_baseline():
     # plt.legend()
     # plt.show()
 
-    np.testing.assert_allclose(actual[:-1], expected[:-1], rtol=0.1)
+    np.testing.assert_allclose(feature.model(t, band, *expected), feature.model(t, band, *actual), rtol=0.1)
 
 
 def test_noisy_all_functions_combination():
@@ -113,8 +113,16 @@ def test_noisy_all_functions_combination():
             # plt.legend()
             # plt.show()
 
+            # The first test might be too rigid. The second test allow for good local minima to be accepted
             np.testing.assert_allclose(actual[:-1], expected[:-1], rtol=0.1)
 
+            # If either the absolute or the relative test passes, it is accepted.
+            # It prevents linexp, which include a flat exactly 0 baseline to not pass the test because
+            # of very minor parameter differences that lead to a major relative difference.
+            np.testing.assert_allclose(
+                feature.model(t, band, *expected), feature.model(t, band, *actual), rtol=0.1, atol=0.1, strict=False
+            )
+
 
 def test_scaler_from_flux_list_input():
     "https://github.com/light-curve/light-curve-python/issues/492"