"""
===========================
Creating a new Source class
===========================
Extending sncosmo with a custom type of Source.
A ``Source`` is something that specifies a spectral timeseries as
a function of an arbitrary number of parameters. For example, the SALT2
model has three parameters (``x0``, ``x1`` and ``c``) that determine a
unique spectrum as a function of phase. The ``SALT2Source`` class implements
the behavior of the model: how the spectral time series depends on those
parameters.
If you have a spectral timeseries model that follows the behavior of one of
the existing classes, such as ``TimeSeriesSource``, great! There's no need to
write a custom class. However, suppose you want to implement a model that
has some new parameterization. In this case, you need a new class that
implements the behavior.
In this example, we implement a new type of source model. Our model is a linear
combination of two spectral time series, with a parameter ``w`` that
determines the relative weight of the models.
"""
import numpy as np
from scipy.interpolate import RectBivariateSpline
import sncosmo
class ComboSource(sncosmo.Source):
_param_names = ['amplitude', 'w']
param_names_latex = ['A', 'w'] # used in plotting display
def __init__(self, phase, wave, flux1, flux2, name=None, version=None):
self.name = name
self.version = version
self._phase = phase
self._wave = wave
# ensure that fluxes are on the same scale
flux2 = flux1.max() / flux2.max() * flux2
self._model_flux1 = RectBivariateSpline(phase, wave, flux1, kx=3, ky=3)
self._model_flux2 = RectBivariateSpline(phase, wave, flux2, kx=3, ky=3)
self._parameters = np.array([1., 0.5]) # initial parameters
def _flux(self, phase, wave):
amplitude, w = self._parameters
return amplitude * ((1.0 - w) * self._model_flux1(phase, wave) +
w * self._model_flux2(phase, wave))
########################################################################
# ... and that's all that we need to define!: A couple class attributes
# (``_param_names`` and ``param_names_latex``, an ``__init__`` method,
# and a ``_flux`` method. The ``_flux`` method is guaranteed to be passed
# numpy arrays for phase and wavelength.
#
# We can now initialize an instance of this source from two spectral time
# series:
#Just as an example, we'll use some undocumented functionality in
# sncosmo to download the Nugent Ia and 2p templates. Don't rely on this
# the `DATADIR` object, or these paths in your code though, as these are
# subject to change between version of sncosmo!
from sncosmo.builtins import DATADIR
phase1, wave1, flux1 = sncosmo.read_griddata_ascii(
DATADIR.abspath('models/nugent/sn1a_flux.v1.2.dat'))
phase2, wave2, flux2 = sncosmo.read_griddata_ascii(
DATADIR.abspath('models/nugent/sn2p_flux.v1.2.dat'))
# In our __init__ method we defined above, the two fluxes need to be on
# the same grid, so interpolate the second onto the first:
flux2_interp = RectBivariateSpline(phase2, wave2, flux2)(phase1, wave1)
source = ComboSource(phase1, wave1, flux1, flux2_interp, name='sn1a+sn2p')
##########################################################################
# We can get a summary of the Source we created:
print(source)
##########################################################################
# Get a spectrum at phase 10 for different parameters:
from matplotlib import pyplot as plt
wave = np.linspace(2000.0, 10000.0, 500)
for w in (0.0, 0.2, 0.4, 0.6, 0.8, 1.0):
source.set(w=w)
plt.plot(wave, source.flux(10., wave), label='w={:3.1f}'.format(w))
plt.legend()
plt.show()
##########################################################################
# The w=0 spectrum is that of the Ia model, the w=1 spectrum is that of
# the IIp model, while intermediate spectra are weighted combinations.
#
# We can even fit the model to some data!
model = sncosmo.Model(source=source)
data = sncosmo.load_example_data()
result, fitted_model = sncosmo.fit_lc(data, model,
['z', 't0', 'amplitude', 'w'],
bounds={'z': (0.2, 1.0),
'w': (0.0, 1.0)})
sncosmo.plot_lc(data, model=fitted_model, errors=result.errors)
##########################################################################
# The fact that the fitted value of w is closer to 0 than 1 indicates that
# the light curve looks more like the Ia template than the IIp template.
# This is generally what we expected since the example data here was
# generated from a Ia template (although not the Nugent template!).