Merge pull request #705 from erusseil/main

UV-extincted BB for Rainbow

Author: erusseil

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

@@ -15,11 +15,8 @@
 __all__ = ["BaseRainbowFit"]
 
 
-# CODATA 2018, grab from astropy
-planck_constant = 6.62607004e-27  # erg s
-speed_of_light = 2.99792458e10  # cm/s
-boltzman_constant = 1.380649e-16  # erg/K
 sigma_sb = 5.6703744191844314e-05  # erg/(cm^2 s K^4)
+speed_of_light = 2.99792458e10  # cm/s
 
 
 IMINUIT_IMPORT_ERROR = (
@@ -69,6 +66,12 @@ def _temperature_parameter_names() -> List[str]:
         """Temperature parameter names."""
         return NotImplementedError
 
+    @staticmethod
+    @abstractmethod
+    def _spectral_parameter_names() -> List[str]:
+        """Spectral model parameter names."""
+        return NotImplementedError
+
     def __post_init__(self) -> None:
         super().__post_init__()
 
@@ -81,6 +84,7 @@ def __post_init__(self) -> None:
             common=self._common_parameter_names(),
             bol=self._bolometric_parameter_names(),
             temp=self._temperature_parameter_names(),
+            spec=self._spectral_parameter_names(),
             bands=self.bands,
             with_baseline=self.with_baseline,
         )
@@ -207,14 +211,6 @@ def _unscale_covariance(self, cov, t_scaler: Scaler, m_scaler: MultiBandScaler)
             cov[:, i] *= scale
             cov[i, :] *= scale
 
-    @staticmethod
-    def planck_nu(wave_cm, T):
-        """Planck function in frequency units."""
-        nu = speed_of_light / wave_cm
-        return (
-            (2 * planck_constant / speed_of_light**2) * nu**3 / np.expm1(planck_constant * nu / (boltzman_constant * T))
-        )
-
     def _lsq_model_no_baseline(self, x, *params):
         """Model function for the fit."""
         t, _band_idx, wave_cm = x
@@ -228,7 +224,8 @@ def _lsq_model_no_baseline(self, x, *params):
         # peak_nu =  2.821 * boltzman_constant * temp / planck_constant # Wien displacement law
         # norm = self.planck_nu(speed_of_light / peak_nu, temp) # Peak = 1 normalization
 
-        planck = self.planck_nu(wave_cm, temp) / norm
+        spectral = self.spectral_func(wave_cm, temp, params)
+        planck = spectral / norm
 
         flux = planck * bol
 

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

@@ -29,6 +29,7 @@ def create_parameters_class(
     common: List[str],
     bol: List[str],
     temp: List[str],
+    spec: List[str],
     bands: Bands,
     with_baseline: bool,
 ):
@@ -37,20 +38,22 @@ def create_parameters_class(
     Parameters
     ----------
     cls_name : str
-        Name of the class to create
+        Name of the class to create.
     common : list of str
-        Common parameters for both bolometric and temperature models
+        Common parameters for both bolometric and temperature models.
     bol : list of str
-        Bolometric model parameters, without common parameters
+        Bolometric model parameters, without common parameters.
     temp : list of str
-        Temperature model parameters, without common parameters
+        Temperature model parameters, without common parameters.
+    spec : list of str
+        Spectral model parameters.
     bands : list of str
         Unique list of bands in the dataset. It is used to generate baseline
         parameters when `with_baseline` is True.
     with_baseline : bool
         Whether to include baseline parameters, one per band in `bands`.
     """
-    attributes = common + bol + temp
+    attributes = common + bol + temp + spec
     if with_baseline:
         baseline = list(map(baseline_parameter_name, bands.names))
         attributes += baseline
@@ -61,14 +64,19 @@ def create_parameters_class(
     enum.common_idx = np.array([enum[attr] for attr in common])
 
     enum.bol = bol
-    enum.bol_idx = np.array([enum[attr] for attr in enum.bol])
+    enum.bol_idx = np.array([enum[attr] for attr in enum.bol], dtype=np.int64)
     enum.all_bol = common + bol
-    enum.all_bol_idx = np.array([enum[attr] for attr in enum.all_bol])
+    enum.all_bol_idx = np.array([enum[attr] for attr in enum.all_bol], dtype=np.int64)
 
     enum.temp = temp
-    enum.temp_idx = np.array([enum[attr] for attr in enum.temp])
+    enum.temp_idx = np.array([enum[attr] for attr in enum.temp], dtype=np.int64)
     enum.all_temp = common + temp
-    enum.all_temp_idx = np.array([enum[attr] for attr in enum.all_temp])
+    enum.all_temp_idx = np.array([enum[attr] for attr in enum.all_temp], dtype=np.int64)
+
+    enum.spec = spec
+    enum.spec_idx = np.array([enum[attr] for attr in enum.spec], dtype=np.int64)
+    enum.all_spec = spec
+    enum.all_spec_idx = np.array([enum[attr] for attr in enum.all_spec], dtype=np.int64)
 
     enum.with_baseline = with_baseline
     if with_baseline:

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

@@ -5,6 +5,7 @@
 from light_curve.light_curve_py.features.rainbow._base import BaseRainbowFit
 
 from .bolometric import BaseBolometricTerm, bolometric_terms
+from .spectral import BaseSpectralTerm, spectral_terms
 from .temperature import BaseTemperatureTerm, temperature_terms
 
 __all__ = ["RainbowFit"]
@@ -59,16 +60,26 @@ class RainbowFit(BaseRainbowFit):
 
     bolometric: Union[str, BaseBolometricTerm] = dataclass_field(default="bazin", kw_only=True)
     """Which parametric bolometric term to use"""
+
     temperature: Union[str, BaseTemperatureTerm] = dataclass_field(default="sigmoid", kw_only=True)
     """Which parametric temperature term to use"""
 
+    spectral: Union[str, BaseSpectralTerm] = dataclass_field(
+        default="planck",
+        kw_only=True,
+    )
+    """Which spectral term to use"""
+
     def __post_init__(self):
         if not isinstance(self.bolometric, BaseBolometricTerm):
             self.bolometric = bolometric_terms[self.bolometric]
 
         if not isinstance(self.temperature, BaseTemperatureTerm):
             self.temperature = temperature_terms[self.temperature]
 
+        if not isinstance(self.spectral, BaseSpectralTerm):
+            self.spectral = spectral_terms[self.spectral]
+
         super().__post_init__()
 
     def _common_parameter_names(self) -> List[str]:
@@ -84,16 +95,22 @@ def _temperature_parameter_names(self) -> List[str]:
         temperature_parameters = self.temperature.parameter_names()
         return [i for i in temperature_parameters if i not in self._common_parameter_names()]
 
+    def _spectral_parameter_names(self) -> List[str]:
+        return self.spectral.parameter_names()
+
     def bol_func(self, t, params):
         return self.bolometric.value(t, *params[self.p.all_bol_idx])
 
     def temp_func(self, t, params):
         return self.temperature.value(t, *params[self.p.all_temp_idx])
 
+    def spectral_func(self, wave_cm, T, params):
+        return self.spectral.value(wave_cm, T, *params[self.p.all_spec_idx])
+
     def _parameter_scalings(self) -> Dict[str, str]:
         rules = super()._parameter_scalings()
 
-        for term in [self.bolometric, self.temperature]:
+        for term in [self.bolometric, self.temperature, self.spectral]:
             for name, scaling in zip(term.parameter_names(), term.parameter_scalings()):
                 rules[name] = scaling
 
@@ -102,12 +119,14 @@ def _parameter_scalings(self) -> Dict[str, str]:
     def _initial_guesses(self, t, m, sigma, band) -> Dict[str, float]:
         initial = self.bolometric.initial_guesses(t, m, sigma, band)
         initial.update(self.temperature.initial_guesses(t, m, sigma, band))
+        initial.update(self.spectral.initial_guesses(t, m, sigma, band))
 
         return initial
 
     def _limits(self, t, m, sigma, band) -> Dict[str, Tuple[float, float]]:
         limits = self.bolometric.limits(t, m, sigma, band)
         limits.update(self.temperature.limits(t, m, sigma, band))
+        limits.update(self.spectral.limits(t, m, sigma, band))
 
         return limits
 

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

@@ -0,0 +1,123 @@
+from abc import abstractmethod
+from dataclasses import dataclass
+from typing import Dict, List, Union
+
+import numpy as np
+
+__all__ = [
+    "spectral_terms",
+    "BaseSpectralTerm",
+    "PlanckSpectralTerm",
+    "BlanketedPlanckSpectralTerm",
+]
+
+# CODATA 2018
+planck_constant = 6.62607004e-27  # erg s
+speed_of_light = 2.99792458e10  # cm/s
+boltzman_constant = 1.380649e-16  # erg/K
+
+
+@dataclass()
+class BaseSpectralTerm:
+    """Spectral term for Rainbow"""
+
+    @staticmethod
+    @abstractmethod
+    def parameter_names() -> List[str]:
+        return NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def parameter_scalings() -> List[Union[str, None]]:
+        return NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def value(wave_cm, T, *params):
+        return NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def initial_guesses(t, m, sigma, band) -> Dict[str, float]:
+        return NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def limits(t, m, sigma, band) -> Dict[str, float]:
+        return NotImplementedError
+
+
+@dataclass()
+class PlanckSpectralTerm(BaseSpectralTerm):
+    """Standard blackbody spectrum"""
+
+    @staticmethod
+    def parameter_names():
+        return []
+
+    @staticmethod
+    def parameter_scalings():
+        return []
+
+    @staticmethod
+    def value(wave_cm, T, *params):
+        nu = speed_of_light / wave_cm
+
+        return (
+            (2 * planck_constant / speed_of_light**2) * nu**3 / np.expm1(planck_constant * nu / (boltzman_constant * T))
+        )
+
+    @staticmethod
+    def initial_guesses(t, m, sigma, band):
+        return {}
+
+    @staticmethod
+    def limits(t, m, sigma, band):
+        return {}
+
+
+@dataclass()
+class BlanketedPlanckSpectralTerm(BaseSpectralTerm):
+    """Blackbody spectrum with exponential blanketing"""
+
+    @staticmethod
+    def parameter_names():
+        return ["log_intensity", "lambda_scale"]
+
+    @staticmethod
+    def parameter_scalings():
+        return [None, None]
+
+    @staticmethod
+    def value(wave_cm, T, log_intensity, lambda_scale):
+        base = PlanckSpectralTerm.value(wave_cm, T)
+
+        # lambda_scale is expressed in Angstrom
+        lambda_scale_cm = lambda_scale * 1e-8
+
+        tau = 10**log_intensity * np.exp(-wave_cm / lambda_scale_cm)
+
+        return base * np.exp(-tau)
+
+    @staticmethod
+    def initial_guesses(t, m, sigma, band):
+        return {
+            "log_intensity": 1.5,
+            "lambda_scale": 100.0,
+        }
+
+    @staticmethod
+    def limits(t, m, sigma, band):
+        # lambda_scale and log_intensity are degenerate. Putting either to ~0 leads to a classical BB.
+        # We prevent log_intensity to go to 0, because the blanketing is more physically realistic
+        # at low lambda_scale high log_intensity rather than high lambda_scale low log_intensity.
+        return {
+            "log_intensity": (1.5, 6.0),
+            "lambda_scale": (100.0, 2000.0),
+        }
+
+
+spectral_terms = {
+    "planck": PlanckSpectralTerm,
+    "blanketed": BlanketedPlanckSpectralTerm,
+}