Merge c59229d into 4a86e45

sncosmo/builtins.py

@@ -27,7 +27,8 @@
     ABMagSystem, CompositeMagSystem, SpectralMagSystem, _MAGSYSTEMS)
 
 from .models import (
-    MLCS2k2Source, SALT2Source, SNEMOSource, TimeSeriesSource, _SOURCES)
+    MLCS2k2Source, SALT2Source, SNEMOSource, SUGARSource,
+    TimeSeriesSource, _SOURCES)
 
 from .specmodel import SpectrumModel
 from .utils import DataMirror
@@ -983,6 +984,24 @@ def load_snemo(relpath, name=None, version=None):
                              version=ver, meta=meta)
 
 
+# SUGAR models
+def load_sugarmodel(relpath, name=None, version=None):
+    abspath = DATADIR.abspath(relpath, isdir=True)
+    return SUGARSource(abspath, name=name, version=version)
+
+
+for name, files, ver in [('sugar', 'sugar', '1.0')]:
+
+    meta = {'type': 'SN Ia', 'subclass': '`~sncosmo.SUGARSource`',
+            'url': 'http://supernovae.in2p3.fr/sugar_template/',
+            'reference': ('Leget20',
+                          'Leget et al. 2020',
+                          '<https://doi.org/10.1051/0004-6361/201834954>')}
+
+    _SOURCES.register_loader(name, load_sugarmodel,
+                             args=['models/sugar/'+files],
+                             version=ver, meta=meta)
+
 # =============================================================================
 # MagSystems
 

sncosmo/models.py

@@ -29,8 +29,8 @@
 from .utils import integration_grid
 
 __all__ = ['get_source', 'Source', 'TimeSeriesSource', 'StretchSource',
-           'SALT2Source', 'MLCS2k2Source', 'SNEMOSource', 'Model',
-           'PropagationEffect', 'CCM89Dust', 'OD94Dust', 'F99Dust']
+           'SUGARSource', 'SALT2Source',  'MLCS2k2Source', 'SNEMOSource',
+           'Model', 'PropagationEffect', 'CCM89Dust', 'OD94Dust', 'F99Dust']
 
 _SOURCES = Registry()
 
@@ -576,6 +576,114 @@ def _flux(self, phase, wave):
                 self._model_flux(phase / self._parameters[1], wave))
 
 
+class SUGARSource(Source):
+    """
+    The SUGAR Type Ia supernova spectral time series template.
+
+    The spectral energy distribution of this model is given by
+
+    .. math::
+
+    F(t, \\lambda) = q_0 10^{-0.4 (M_0(t, \\lambda)
+                             + q_1 \alpha_1(t, \\lambda)
+                             + q_2 \alpha_2(t, \\lambda)
+                             + q_3 \alpha_3(t, \\lambda)
+                             + A_v CCM(\\lambda))}
+                             (10^{-3} c\\lambda^{2})
+
+    where ``q_0``, ``q_1``, ``q_2``, ``q_3``,
+    and ``A_v`` are the free parameters
+    of the model,``alpha_0``, ``alpha_1``,
+    `alpha_2``, `alpha_3``, `CCM`` are
+    the template vectors of the model.
+    The ``q_0`` is the equivalent parameter
+    in flux of the ``Delta M_{gray}``
+    parameter define in Leget et al. 2020.
+
+    Parameters
+    ----------
+    modeldir : str, optional
+        Directory path containing model component files. Default is `None`,
+        which means that no directory is prepended to filenames when
+        determining their path.
+
+    m0file : str or fileobj, optional
+    alpha1file : str or fileobj, optional
+    alpha2file : str or fileobj, optional
+    alpha3file : str or fileobj, optional
+    CCMfile: str or fileobj, optional
+        Filenames of various model components. Defaults are:
+        * m0file = 'sugar_template_0.dat' (2-d grid)
+        * alpha1file = 'sugar_template_1.dat' (2-d grid)
+        * alpha2file = 'sugar_template_2.dat' (2-d grid)
+        * alpha3file = 'sugar_template_3.dat' (2-d grid)
+        * CCMfile = 'sugar_template_4.dat' (2-d grid)
+
+    Notes
+    -----
+    The "2-d grid" files have the format ``<phase> <wavelength>
+    <value>`` on each line.
+    """
+    _param_names = ['q0', 'q1', 'q2', 'q3', 'Av']
+    param_names_latex = ['q_0', 'q_1', 'q_2', 'q_3', 'A_v']
+
+    def __init__(self, modeldir=None,
+                 m0file='sugar_template_0.dat',
+                 alpha1file='sugar_template_1.dat',
+                 alpha2file='sugar_template_2.dat',
+                 alpha3file='sugar_template_3.dat',
+                 CCMfile='sugar_template_4.dat',
+                 name=None, version=None):
+
+        self.name = name
+        self.version = version
+        self._model = {}
+        self.M_keys = ['M0', 'ALPHA1', 'ALPHA2', 'ALPHA3', 'CCM']
+        self._parameters = np.zeros(len(self.M_keys))
+        self._parameters[0] = 1e-15
+        names_or_objs = {'M0': m0file,
+                         'ALPHA1': alpha1file,
+                         'ALPHA2': alpha2file,
+                         'ALPHA3': alpha3file,
+                         'CCM': CCMfile}
+
+        # Make filenames into full paths.
+        if modeldir is not None:
+            for k in names_or_objs:
+                v = names_or_objs[k]
+                if (v is not None and isinstance(v, str)):
+                    names_or_objs[k] = os.path.join(modeldir, v)
+
+        for i, key in enumerate(self.M_keys):
+            phase, wave, values = read_griddata_ascii(names_or_objs[key])
+            self._model[key] = BicubicInterpolator(phase, wave, values)
+            if key == 'M0':
+                # The "native" phases and wavelengths of the model are those
+                self._phase = np.linspace(-12., 48, 21)
+                self._wave = wave
+
+    def _flux(self, phase, wave):
+        mag_sugar = self._model['M0'](phase, wave)
+        for i, key in enumerate(self.M_keys):
+            if key != 'M0':
+                comp = self._model[key](phase, wave) * self._parameters[i]
+                mag_sugar += comp
+        # Mag AB used in the training of SUGAR.
+        mag_sugar += 48.59
+        wave_factor = (wave ** 2 / 299792458. * 1.e-10)
+        flux = (self._parameters[0] * 10. ** (-0.4 * mag_sugar) / wave_factor)
+
+        if hasattr(phase, '__iter__'):
+            not_define = ~((phase>-12) & (phase<48))
+            flux[not_define] = 0
+            return flux
+        else:
+            if phase<-12 or phase>48:
+                return np.zeros_like(wave)
+            else:
+                return flux
+
+
 class SALT2Source(Source):
     """The SALT2 Type Ia supernova spectral timeseries model.
 

sncosmo/tests/test_sugarsource.py

@@ -0,0 +1,68 @@
+# Licensed under a 3-clause BSD style license - see LICENSES
+
+"""Tests for SUGARSource (and wrapped in Model)"""
+
+import os
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+import sncosmo
+
+
+def sugar_model(q0=0, q1=0, q2=0, q3=0, Av=0,
+                phase=np.linspace(-5, 30, 10),
+                wave=np.linspace(4000, 8000, 10)):
+    """
+    Give a spectral time series of SUGAR model
+    for a given set of parameters.
+    """
+    source = sncosmo.get_source('sugar')
+
+    mag_sugar = source._model['M0'](phase, wave)
+
+    keys = ['ALPHA1', 'ALPHA2', 'ALPHA3', 'CCM']
+    parameters = [q1, q2, q3, Av]
+    for i in range(4):
+        comp = source._model[keys[i]](phase, wave) * parameters[i]
+        mag_sugar += comp
+    # Mag AB used in the training of SUGAR.
+    mag_sugar += 48.59
+    wave_factor = (wave ** 2 / 299792458. * 1.e-10)
+    return (q0 * 10. ** (-0.4 * mag_sugar) / wave_factor)
+
+
+@pytest.mark.might_download
+def test_sugarsource():
+    """Test timeseries output from SUGARSource vs pregenerated timeseries
+    from the original files."""
+
+    source = sncosmo.get_source("sugar")
+    model = sncosmo.Model(source)
+
+    q1 = [-1, 0, 1, 2]
+    q2 = [1, 0, -1, -2]
+    q3 = [-1, 1, 0, -2]
+    Av = [-0.1, 0, 0.2, 0.5]
+    q0 = [10**(-0.4 * 34), 10**(-0.4 * 33),
+          10**(-0.4 * 38), 10**(-0.4 * 42)]
+
+    time = np.linspace(-5, 30, 10)
+    wave = np.linspace(4000, 8000, 10)
+
+    for i in range(len(q1)):
+
+        fluxref = sugar_model(q0=q0[i],
+                              q1=q1[i],
+                              q2=q2[i],
+                              q3=q3[i],
+                              Av=Av[i],
+                              phase=time,
+                              wave=wave)
+
+        model.set(z=0, t0=0, q0=q0[i],
+                  q1=q1[i], q2=q2[i],
+                  q3=q3[i], Av=Av[i])
+        flux = model.flux(time, wave)
+        assert_allclose(flux, fluxref, rtol=1e-13)