Add phase dependence to prop. effects (#359)

sncosmo/models.py

@@ -1546,26 +1546,31 @@ def _flux(self, time, wave):
         """Array flux function."""
 
         a = 1. / (1. + self._parameters[0])
-        phase = (time - self._parameters[1]) * a
+        obsphase = (time - self._parameters[1])
+        restphase = obsphase * a
         restwave = wave * a
 
         # Note that below we multiply by the scale factor to conserve
         # bolometric luminosity.
-        f = a * self._source._flux(phase, restwave)
+        f = a * self._source._flux(restphase, restwave)
 
         # Pass the flux through the PropagationEffects.
         for effect, frame, zindex in zip(self._effects, self._effect_frames,
                                          self._effect_zindicies):
             if frame == 'obs':
                 effect_wave = wave
+                effect_phase = obsphase
             elif frame == 'rest':
                 effect_wave = restwave
+                effect_phase = restphase
             else:  # frame == 'free'
                 effect_a = 1. / (1. + self._parameters[zindex])
                 effect_wave = wave * effect_a
-
-            f = effect.propagate(effect_wave, f)
-
+                effect_phase = obsphase * effect_a
+            try:
+                f = effect.propagate(effect_wave, f, phase=effect_phase)
+            except TypeError:
+                f = effect.propagate(effect_wave, f)
         return f
 
     def flux(self, time, wave):
@@ -1948,15 +1953,30 @@ def minwave(self):
     def maxwave(self):
         return self._maxwave
 
+    def minphase(self):
+        try:
+            return self._minphase
+        except AttributeError:
+            return np.nan
+
+    def maxphase(self):
+        try:
+            return self._maxphase
+        except AttributeError:
+            return np.nan
+
     @abc.abstractmethod
-    def propagate(self, wave, flux):
+    def propagate(self, wave, flux, phase=None):
         pass
 
     def _headsummary(self):
         summary = """\
         class           : {0}
-        wavelength range: [{1:.6g}, {2:.6g}] Angstroms"""\
-        .format(self.__class__.__name__, self._minwave, self._maxwave)
+        wavelength range: [{1:.6g}, {2:.6g}] Angstroms
+        phase range     : [{3:.2g}, {4:.2g}]"""\
+        .format(self.__class__.__name__,
+                self._minwave, self._maxwave,
+                self.minphase(), self.maxphase())
         return dedent(summary)
 
 
@@ -1970,7 +1990,7 @@ class CCM89Dust(PropagationEffect):
     def __init__(self):
         self._parameters = np.array([0., 3.1])
 
-    def propagate(self, wave, flux):
+    def propagate(self, wave, flux, phase=None):
         """Propagate the flux."""
         ebv, r_v = self._parameters
         return extinction.apply(extinction.ccm89(wave, ebv * r_v, r_v), flux)
@@ -1986,7 +2006,7 @@ class OD94Dust(PropagationEffect):
     def __init__(self):
         self._parameters = np.array([0., 3.1])
 
-    def propagate(self, wave, flux):
+    def propagate(self, wave, flux, phase=None):
         """Propagate the flux."""
         ebv, r_v = self._parameters
         return extinction.apply(extinction.odonnell94(wave, ebv * r_v, r_v),
@@ -2005,7 +2025,7 @@ def __init__(self, r_v=3.1):
         self._r_v = r_v
         self._f = extinction.Fitzpatrick99(r_v=r_v)
 
-    def propagate(self, wave, flux):
+    def propagate(self, wave, flux, phase=None):
         """Propagate the flux."""
         ebv = self._parameters[0]
         return extinction.apply(self._f(wave, ebv * self._r_v), flux)

sncosmo/tests/test_models.py

@@ -2,6 +2,8 @@
 
 from io import StringIO
 
+import scipy
+
 import numpy as np
 from numpy.testing import assert_allclose, assert_approx_equal
 
@@ -18,9 +20,48 @@ def flatsource():
     return sncosmo.TimeSeriesSource(phase, wave, flux)
 
 
+class AchromaticMicrolensing(sncosmo.PropagationEffect):
+    """
+    An achromatic microlensing object. Tests phase dependence.
+    """
+    _param_names = []
+    param_names_latex = []
+    _minwave = np.nan
+    _maxwave = np.nan
+    _minphase = np.nan
+    _maxphase = np.nan
+
+    def __init__(self, time, magnification):
+        """
+        Parameters
+            ----------
+            time: :class:`~list` or :class:`~numpy.array`
+                A time array for your microlensing
+            dmag: :class:`~list` or :class:`~numpy.array`
+                microlensing magnification
+        """
+        self._parameters = np.array([])
+        self.mu = scipy.interpolate.interp1d(time,
+                                             magnification,
+                                             bounds_error=False,
+                                             fill_value=1.)
+        self._minphase = np.min(time)
+        self._maxphase = np.max(time)
+
+    def propagate(self, wave, flux, phase):
+        """
+        Propagate the magnification onto the model's flux output.
+        """
+
+        mu = np.expand_dims(np.atleast_1d(self.mu(phase)), 1)
+        return flux * mu
+
+
 class StepEffect(sncosmo.PropagationEffect):
-    """Effect with transmission 0 below cutoff wavelength, 1 above.
-    Useful for testing behavior with redshift."""
+    """
+    Effect with transmission 0 below cutoff wavelength, 1 above.
+    Useful for testing behavior with redshift.
+    """
 
     _param_names = ['stepwave']
     param_names_latex = [r'\lambda_{s}']
@@ -312,3 +353,17 @@ def test_effect_frame_free():
 
     wave = np.array([12000., 13000., 14000., 14999., 15000., 16000.])
     assert_allclose(model.flux(0., wave), [0., 0., 0., 0., 0.5, 0.5])
+
+
+def test_effect_phase_dependent():
+    """Test Model with PropagationEffect with phase dependence"""
+
+    model = sncosmo.Model(source=flatsource())
+    flux = model.flux(50., 5000.).flatten()
+
+    model_micro = sncosmo.Model(source=flatsource(),
+                                effects=[AchromaticMicrolensing([40., 60.],
+                                                                [10., 10.])],
+                                effect_frames=['rest'],
+                                effect_names=['micro'])
+    assert_approx_equal(model_micro.flux(50., 5000.).flatten(), flux*10.)