%pylab inline

import sys
from corner import corner
sys.path.append('.')

from src.mocklc import MockLC, SimulationSetup
from src.blendlpf import MockLPF
import src.plotting as pl 

"""
Explanation: Contamination example 2a
M star slightly contaminated by a G star and uninformative prior on orbital parameters
Hannu Parviainen<br>
Instituto de Astrofísica de Canarias
Last modified: 15.7.2019
Here we use the pytransit.contamination module to estimate the true planet to star radius ratio robustly using multicolour photometry in the presence of possible flux contamination from an unresolved source in the photometry aperture, as detailed in Parviainen et al. 2019 (submitted). This can be used in the validation of transiting planet candidates, where, e.g., blended eclipsing binaries are a significant source of false positives.
This notebook runs the simulation for a light curve without contamination and informative priors on the impact parameter and stellar density, while the previous notebook (1a) did the same with uninformative priors on the orbital parameters.
Light curves: We don't use real data here, but create simulated multicolour photometry lightcurves using the MockLC class found in src.mocklc. The code is the same that was used for the simulations in Parviainen et al. (2019).
Log posterior function: The log posterior function is defined by MockLPF class found in src.blendlpf.MockLPF. The class inherits pytransit.lpf.PhysContLPF and overrides the _init_instrument method to define the instrument and the contamination model (amongst other things to make running a variety of simulations smooth).
Parametrisation: As discussed in the paper, the contamination is parametrised by the apparent area ratio ($k_\mathrm{True}^2$), true area ratio ($k_\mathrm{App}^2$), and the effective temperatures of the host and contaminant stars. The apparent area ratio defines how deep the transit is in a single single passband and can be wavelength dependent (if the host and contaminant are of different spectral type), while the true area ratio stands for the unblended true geometric planet-star area ratio.
The true radius ratio ($k_\mathrm{True}$) is the main quantity of interest in transiting planet candidate validation), since it together with a stellar radius estimate gives the true absolute planetary radius.
End of explanation
"""


lc = MockLC(SimulationSetup('M', 0.1, 0.0, 0.15, 'short_transit', cteff=5500, know_orbit=False))
lc.create(wnsigma=[0.001, 0.001, 0.001, 0.001], rnsigma=0.00001, rntscale=0.5, nights=1);
lc.plot();

"""
Explanation: Create a mock light curve
End of explanation
"""


lpf = MockLPF('Example_1', lc)

lpf.print_parameters(columns=2)

"""
Explanation: Initialize the log posterior function
End of explanation
"""


lpf.optimize_global(1000)

lpf.plot_light_curves()

"""
Explanation: Optimize
End of explanation
"""


lpf.sample_mcmc(5000, reset=True, repeats=2)

"""
Explanation: Estimate the posterior
The contamination parameter space is a bit difficult to sample (especially if the signal to noise ratio is low), so the sampling should be continued at least for 10000 iterations.
End of explanation
"""


df = lpf.posterior_samples()

pl.joint_radius_ratio_plot(df, fw=13, clim=(0.099, 0.12), htelim=(3570, 3630), ctelim=(4000,7500), blim=(0, 0.5), rlim=(3.8, 5.2));

pl.joint_contamination_plot(df, fw=13, clim=(0, 0.4), htelim=(3570, 3630), ctelim=(4000,7500), blim=(0, 0.5), rlim=(3.8, 5.2));

"""
Explanation: Analysis
We plot the main results below. It is clear that a single good-quality four-colour light curve still allows for contamination from a source of similar spectral type as the host star. However, in this example, the maximum allowed level of contamination is not sufficient to take the transiting object out of the planetary regime.
Also, the joint posterior plots clearly show that any significant contamination must come from a source of a similar spectral type as the host. Combining this information with prior knowledge about the probability of having such a system without colour variations can be used in probabilistic planet candidate validation.
Plot the basic joint posterior
End of explanation
"""


pl.marginal_radius_ratio_plot(df, bins=60, klim=(0.097, 0.12), figsize=(7,5));

"""
Explanation: Plot the apparent and true radius ratio posteriors
End of explanation
"""


corner(df.iloc[:,2:-3]);

"""
Explanation: Make a corner plot to have a good overview to the posterior space
End of explanation
"""


