Merge pull request #407 from karpov-sv/rainbow_ml

Maximum likelihood for Rainbow

Author: hombit

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

@@ -1,5 +1,4 @@
 from abc import abstractmethod
-from copy import deepcopy
 from dataclasses import dataclass
 from typing import Dict, List, Tuple
 
@@ -11,6 +10,7 @@
 from light_curve.light_curve_py.features.rainbow._parameters import create_parameters_class
 from light_curve.light_curve_py.features.rainbow._scaler import MultiBandScaler, Scaler
 from light_curve.light_curve_py.minuit_lsq import LeastSquares
+from light_curve.light_curve_py.minuit_ml import MaximumLikelihood
 
 __all__ = ["BaseRainbowFit"]
 
@@ -121,6 +121,9 @@ def _check_iminuit():
         if LeastSquares is None:
             raise ImportError(IMINUIT_IMPORT_ERROR)
 
+        if MaximumLikelihood is None:
+            raise ImportError(IMINUIT_IMPORT_ERROR)
+
         try:
             try:
                 from packaging.version import parse as parse_version
@@ -144,48 +147,65 @@ def temp_func(self, t, params):
         """Temperature evolution function."""
         return NotImplementedError
 
-    @abstractmethod
-    def _unscale_parameters(self, params, t_scaler: Scaler, m_scaler: MultiBandScaler) -> None:
-        """Unscale parameters from internal units, in-place.
+    def _parameter_scalings(self) -> Dict[str, str]:
+        """Rules for scaling/unscaling the parameters"""
+        rules = {}
 
-        No baseline parameters are needed to be unscaled.
-        """
-        return NotImplementedError
+        if self.with_baseline:
+            for band_name in self.bands.names:
+                baseline_name = self.p.baseline_parameter_name(band_name)
+                rules[baseline_name] = "baseline"
 
-    def _unscale_errors(self, errors, t_scaler: Scaler, m_scaler: MultiBandScaler) -> None:
-        """Unscale parameter errors from internal units, in-place.
+        return rules
 
-        No baseline parameters are needed to be unscaled.
-        """
+    def _parameter_scale(self, name: str, t_scaler: Scaler, m_scaler: MultiBandScaler) -> float:
+        """Return the scale factor to be applied to the parameter to unscale it"""
+        scaling = self._parameter_scalings().get(name)
+        if scaling == "time" or scaling == "timescale":
+            return t_scaler.scale
+        elif scaling == "flux":
+            return m_scaler.scale
 
-        # We need to modify original scalers to only apply the scale, not shifts, to the errors
-        # It should be re-implemented in subclasses for a cleaner way to unscale the errors
-        t_scaler = deepcopy(t_scaler)
-        m_scaler = deepcopy(m_scaler)
-        t_scaler.reset_shift()
-        m_scaler.reset_shift()
+        return 1
 
-        return self._unscale_parameters(errors, t_scaler, m_scaler)
+    def _unscale_parameters(self, params, t_scaler: Scaler, m_scaler: MultiBandScaler) -> None:
+        """Unscale parameters from internal units, in-place."""
+        for name, scaling in self._parameter_scalings().items():
+            if scaling == "time":
+                params[self.p[name]] = t_scaler.undo_shift_scale(params[self.p[name]])
 
-    def _unscale_baseline_parameters(self, params, m_scaler: MultiBandScaler) -> None:
-        """Unscale baseline parameters from internal units, in-place.
+            elif scaling == "timescale":
+                params[self.p[name]] = t_scaler.undo_scale(params[self.p[name]])
 
-        Must be used only if `with_baseline` is True.
-        """
-        for band_name in self.bands.names:
-            baseline_name = self.p.baseline_parameter_name(band_name)
-            baseline = params[self.p[baseline_name]]
-            params[self.p[baseline_name]] = m_scaler.undo_shift_scale_band(baseline, band_name)
+            elif scaling == "flux":
+                params[self.p[name]] = m_scaler.undo_scale(params[self.p[name]])
 
-    def _unscale_baseline_errors(self, errors, m_scaler: MultiBandScaler) -> None:
-        """Unscale baseline parameters from internal units, in-place.
+            elif scaling == "baseline":
+                band_name = self.p.baseline_band_name(name)
+                baseline = params[self.p[name]]
+                params[self.p[name]] = m_scaler.undo_shift_scale_band(baseline, band_name)
 
-        Must be used only if `with_baseline` is True.
-        """
-        for band_name in self.bands.names:
-            baseline_name = self.p.baseline_parameter_name(band_name)
-            baseline = errors[self.p[baseline_name]]
-            errors[self.p[baseline_name]] = m_scaler.undo_scale_band(baseline, band_name)
+                pass
+
+            elif scaling is None or scaling.lower() == "none":
+                pass
+
+            else:
+                raise ValueError("Unsupported parameter scaling: " + scaling)
+
+    def _unscale_errors(self, errors, t_scaler: Scaler, m_scaler: MultiBandScaler) -> None:
+        """Unscale parameter errors from internal units, in-place."""
+        for name in self.names:
+            scale = self._parameter_scale(name, t_scaler, m_scaler)
+            errors[self.p[name]] *= scale
+
+    def _unscale_covariance(self, cov, t_scaler: Scaler, m_scaler: MultiBandScaler) -> None:
+        """Unscale parameter covariance from internal units, in-place."""
+        for name in self.names:
+            scale = self._parameter_scale(name, t_scaler, m_scaler)
+            i = self.p[name]
+            cov[:, i] *= scale
+            cov[i, :] *= scale
 
     @staticmethod
     def planck_nu(wave_cm, T):
@@ -283,7 +303,19 @@ def _eval(self, *, t, m, sigma, band):
     def _eval_and_fill(self, *, t, m, sigma, band, fill_value):
         return super()._eval_and_fill(t=t, m=m, sigma=sigma, band=band, fill_value=fill_value)
 
-    def _eval_and_get_errors(self, *, t, m, sigma, band, print_level=None, get_initial=False):
+    def _eval_and_get_errors(
+        self,
+        *,
+        t,
+        m,
+        sigma,
+        band,
+        upper_mask=None,
+        get_initial=False,
+        return_covariance=False,
+        print_level=None,
+        debug=False,
+    ):
         # Initialize data scalers
         t_scaler = Scaler.from_time(t)
         m_scaler = MultiBandScaler.from_flux(m, band, with_baseline=self.with_baseline)
@@ -311,19 +343,51 @@ def _eval_and_get_errors(self, *, t, m, sigma, band, print_level=None, get_initi
             initial_guesses = self._initial_guesses(t, m, sigma, band)
             limits = self._limits(t, m, sigma, band)
 
-        least_squares = LeastSquares(
+        # least_squares = LeastSquares(
+        cost_function = MaximumLikelihood(
             model=self._lsq_model,
             parameters=limits,
             x=(t, band_idx, wave_cm),
             y=m,
             yerror=sigma,
+            upper_mask=upper_mask,
         )
-        minuit = self.Minuit(least_squares, name=self.names, **initial_guesses)
+        minuit = self.Minuit(cost_function, name=self.names, **initial_guesses)
         # TODO: expose these parameters through function arguments
         if print_level is not None:
             minuit.print_level = print_level
-        minuit.strategy = 2
-        minuit.migrad(ncall=10000, iterate=10)
+        minuit.strategy = 0  # We will need to manually call .hesse() on convergence anyway
+
+        # Supposedly it is not the same as just setting iterate=10?..
+        for i in range(10):
+            minuit.migrad()
+
+            if minuit.valid:
+                minuit.hesse()
+                # hesse() may may drive it invalid
+                if minuit.valid:
+                    break
+            else:
+                # That's what iterate is supposed to do?..
+                minuit.simplex()
+                # FIXME: it may drive the fit valid, but we will not have Hesse run on last iteration
+
+        if debug:
+            # Expose everything we have to outside, unscaled, for easier debugging
+            self.minuit = minuit
+            self.mparams = {
+                "t": t,
+                "band_idx": band_idx,
+                "wave_cm": wave_cm,
+                "m": m,
+                "sigma": sigma,
+                "limits": limits,
+                "upper_mask": upper_mask,
+                "initial_guesses": initial_guesses,
+                "values": minuit.values,
+                "errors": minuit.errors,
+                "covariance": minuit.covariance,
+            }
 
         if not minuit.valid and self.fail_on_divergence and not get_initial:
             raise RuntimeError("Fitting failed")
@@ -338,15 +402,19 @@ def _eval_and_get_errors(self, *, t, m, sigma, band, print_level=None, get_initi
         errors = np.array(minuit.errors)
 
         self._unscale_parameters(params, t_scaler, m_scaler)
-        if self.with_baseline:
-            self._unscale_baseline_parameters(params, m_scaler)
 
         # Unscale errors
         self._unscale_errors(errors, t_scaler, m_scaler)
-        if self.with_baseline:
-            self._unscale_baseline_errors(errors, m_scaler)
 
-        return np.r_[params, reduced_chi2], errors
+        return_values = np.r_[params, reduced_chi2], errors
+
+        if return_covariance:
+            # Unscale covaiance
+            cov = np.array(minuit.covariance)
+            self._unscale_covariance(cov, t_scaler, m_scaler)
+            return_values += (cov,)
+
+        return return_values
 
     def fit_and_get_errors(self, t, m, sigma, band, *, sorted=None, check=True, **kwargs):
         t, m, sigma, band = self._normalize_input(t=t, m=m, sigma=sigma, band=band, sorted=sorted, check=check)

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

@@ -12,6 +12,13 @@ def baseline_parameter_name(band: str) -> str:
     return f"baseline_{band}"
 
 
+def baseline_band_name(name: str) -> str:
+    if name.startswith("baseline_"):
+        return name[len("baseline_") :]
+
+    return None
+
+
 def create_int_enum(cls_name: str, attributes: Iterable[str]):
     return IntEnum(cls_name, {attr: i for i, attr in enumerate(attributes)})
 
@@ -68,6 +75,7 @@ def create_parameters_class(
         enum.all_baseline = baseline
         enum.baseline_idx = np.array([enum[attr] for attr in enum.all_baseline])
         enum.baseline_parameter_name = staticmethod(baseline_parameter_name)
+        enum.baseline_band_name = staticmethod(baseline_band_name)
 
         band_idx_to_baseline_idx = {
             band_idx: enum[baseline_parameter_name(band_name)] for band_idx, band_name in zip(bands.index, bands.names)

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

@@ -47,10 +47,6 @@ def do_scale(self, x):
     def undo_scale(self, x):
         return x * self.scale
 
-    def reset_shift(self):
-        """Resets scaler shift to zero, keeping only the scale"""
-        self.shift *= 0
-
 
 @dataclass()
 class MultiBandScaler(Scaler):
@@ -59,9 +55,6 @@ class MultiBandScaler(Scaler):
     per_band_shift: Dict[str, float]
     """Shift to apply to each band"""
 
-    per_band_scale: Dict[str, float]
-    """Scale to apply to each band"""
-
     @classmethod
     def from_flux(cls, flux, band, *, with_baseline: bool) -> "MultiBandScaler":
         """Create a Scaler from a flux array.
@@ -71,7 +64,7 @@ def from_flux(cls, flux, band, *, with_baseline: bool) -> "MultiBandScaler":
         """
         uniq_bands = np.unique(band)
         per_band_shift = dict.fromkeys(uniq_bands, 0.0)
-        shift_array = np.zeros_like(flux)
+        shift_array = np.zeros(len(flux))
 
         if with_baseline:
             for b in uniq_bands:
@@ -81,19 +74,8 @@ def from_flux(cls, flux, band, *, with_baseline: bool) -> "MultiBandScaler":
         scale = np.std(flux)
         if scale == 0.0:
             scale = 1.0
-        per_band_scale = dict.fromkeys(uniq_bands, scale)
 
-        return cls(shift=shift_array, scale=scale, per_band_shift=per_band_shift, per_band_scale=per_band_scale)
+        return cls(shift=shift_array, scale=scale, per_band_shift=per_band_shift)
 
     def undo_shift_scale_band(self, x, band):
-        return x * self.per_band_scale.get(band, 1) + self.per_band_shift.get(band, 0)
-
-    def undo_scale_band(self, x, band):
-        return x * self.per_band_scale.get(band, 1)
-
-    def reset_shift(self):
-        """Resets scaler shift to zero, keeping only the scale"""
-        for band in self.per_band_shift:
-            self.per_band_shift[band] = 0
-
-        super().reset_shift()
+        return x * self.scale + self.per_band_shift.get(band, 0)

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

@@ -71,7 +71,7 @@ def parameter_scalings():
     def value(t, t0, amplitude, rise_time):
         dt = t - t0
 
-        result = np.zeros_like(dt)
+        result = np.zeros(len(dt))
         # To avoid numerical overflows, let's only compute the exponents not too far from t0
         idx = dt > -100 * rise_time
         result[idx] = amplitude / (np.exp(-dt[idx] / rise_time) + 1)
@@ -80,7 +80,7 @@ def value(t, t0, amplitude, rise_time):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        A = np.max(m)
+        A = np.ptp(m)
 
         initial = {}
         initial["reference_time"] = t[np.argmax(m)]
@@ -92,12 +92,14 @@ def initial_guesses(t, m, sigma, band):
     @staticmethod
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
-        m_amplitude = np.max(m)
+        m_amplitude = np.ptp(m)
+
+        mean_dt = np.median(t[1:] - t[:-1])
 
         limits = {}
         limits["reference_time"] = (np.min(t) - 10 * t_amplitude, np.max(t) + 10 * t_amplitude)
-        limits["amplitude"] = (0.0, 10 * m_amplitude)
-        limits["rise_time"] = (1e-4, 10 * t_amplitude)
+        limits["amplitude"] = (0.0, 20 * m_amplitude)
+        limits["rise_time"] = (0.1 * mean_dt, 10 * t_amplitude)
 
         return limits
 
@@ -128,7 +130,7 @@ def value(t, t0, amplitude, rise_time, fall_time):
             -fall_time / (fall_time + rise_time)
         )
 
-        result = np.zeros_like(dt)
+        result = np.zeros(len(dt))
         # To avoid numerical overflows, let's only compute the exponents not too far from t0
         idx = (dt > -100 * rise_time) & (dt < 100 * fall_time)
         result[idx] = amplitude * scale / (np.exp(-dt[idx] / rise_time) + np.exp(dt[idx] / fall_time))
@@ -137,15 +139,17 @@ def value(t, t0, amplitude, rise_time, fall_time):
 
     @staticmethod
     def initial_guesses(t, m, sigma, band):
-        A = np.max(m)
+        A = np.ptp(m)
+
+        mc = m - np.min(m)  # To avoid crashing on all-negative data
 
         # Naive peak position from the highest point
         t0 = t[np.argmax(m)]
-        # Peak position as weighted centroid of everything above zero
-        idx = m > 0
+        # 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 * m[idx] / sigma[idx]) / np.sum(m[idx] / sigma[idx]))
+        dt = np.sqrt(np.sum((t[idx] - t0) ** 2 * (mc[idx]) / sigma[idx]) / np.sum(mc[idx] / sigma[idx]))
 
         # Empirical conversion of sigma to rise/fall times
         rise_time = dt / 2
@@ -165,13 +169,15 @@ def initial_guesses(t, m, sigma, band):
     @staticmethod
     def limits(t, m, sigma, band):
         t_amplitude = np.ptp(t)
-        m_amplitude = np.max(m)
+        m_amplitude = np.ptp(m)
+
+        mean_dt = np.median(t[1:] - t[:-1])
 
         limits = {}
         limits["reference_time"] = (np.min(t) - 10 * t_amplitude, np.max(t) + 10 * t_amplitude)
-        limits["amplitude"] = (0.0, 10 * m_amplitude)
-        limits["rise_time"] = (1e-4, 10 * t_amplitude)
-        limits["fall_time"] = (1e-4, 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)
 
         return limits