%pylab inline

"""
Explanation: Flux contamination 1
Black body model
Hannu Parviainen<br>
Edited: 1.1.2020
Introduction
The contamination model is aimed for simulating exoplanet transit light curves where the planet host star flux is "contaminated" by flux from nearby stars (also called blending). Contamination leads to a decrease in the observed transit depth (the planet appears to be smaller than it truly is), and this effect is achromatic if the host and the contaminant(s) are of different spectral type.
The black body contamination model pytransit.contamination.BBContamination approximates the host and the contaminant star(s) as black bodies defined by their effective temperature $T_\mathrm{Eff}$. This simplification can be useful, but pytransit.contamination.SMContamination offers a more physically accurate model.
End of explanation
"""


from pytransit.contamination import BBContamination, Instrument, sdss_griz

"""
Explanation: Import
End of explanation
"""


cn = BBContamination(Instrument('example', sdss_griz), ref_pb="r'")

"""
Explanation: Initialisation
The model is initialised with a pytransit.contamination.Instrument that defines the passbands (filters). We'll set up a simple instrument with approzimate sloan g', r', i', and z filters available as a set as pytransit.contamination.sdss_griz. The initialisation of BBContaminationalso requires the name of the reference passband, ref_pb, that has to correspond to one of the filter names in the instrument.
End of explanation
"""


cn.contamination(0.25, 4400, 6200)

"""
Explanation: Usage
After the initialisation, the contamination ratio can be calculated using BBContamination.contamination, that takes the contamination in the reference passband (cref), the effective temperature of the host (teff1), and the effective temperature of the contaminant (teff2).
End of explanation
"""


fig = cn.plot(3500, 5800, 0.25, nwl=1000, figsize=(13,8));

fig, axs = subplots(2, 2, figsize=(13,6), constrained_layout=True)
cn.plot(3400, 6200, 0.25, wlref=630, axs=axs[:,0])
cn.plot(6200, 3400, 0.25, wlref=630, axs=axs[:,1])

"""
Explanation: The model can also be plotted as a function of wavelength using BBContamination.plot.
End of explanation
"""


from pytransit import UniformModel
from pytransit.contamination import contaminate_light_curve

tm = UniformModel()

npt = 500
npb = len(cn.instrument.filters)
times = tile(linspace(0.88, 1.12, npt), npb)
lcids = repeat(arange(npb), npt)
pbids = arange(npb)

tm.set_data(times, lcids, pbids)
flux = tm.evaluate_ps(0.1, 1.0, 1.5, 4.0, 0.5*pi)

c = cn.contamination(0.5, 4400, 6200)
cflux = contaminate_light_curve(flux, c, pbids[lcids])

fig, ax = subplots(figsize=(13,6))
ax.plot(flux)
ax.plot(cflux)
fig.tight_layout()

"""
Explanation: Transit light curve example
We use a pytransit.UniformModel so we don't need to think about the limb darkening.
The transit model needs an array of mid-exposure times (times), light curve indices (lcids), and passband indices (pbids). The lcids and pbids indices are needed since PyTransit models heterogeneous light curves consisting of  multiple passbands: lcids map each exposure to a single light curve, and pbids map each light curve to a single passband.
End of explanation
"""


