%pylab inline

from pytransit import RoadRunnerModel

rc('figure', figsize=(13,4))

tm = RoadRunnerModel()

"""
Explanation: TTV model evaluation example
Author: Hannu Parviainen<br/>
Last modified: 20.10.2020
Evaluating the transit model for a planet whose transit center times deviate from a linear ephemeris is extremely simple in PyTransit. Each light curve (indexed by an array of light curve indices, lcids) can be given its own zero epoch by giving the model a set of per-light-curve epoch indices (epids), and then giving the zero epoch (t0) as an array rather than as a single float. Here is an example using the RoadRunner model, but the same approach will work with the QuadraticModel as well (and with the rest of the transit models a bit later).
Note: The first evaluation of the transit model takes a bit of time because numba compiles the code behind the scenes only when needed (and calling TransitModel.evaluate with different types of parameters calls actually a different numba routines). The model evaluation is extremely fast after all the required code has been compiled.
End of explanation
"""


ntr = 3
npt = 200

times = concatenate(linspace(-0.2, 0.2, npt)[newaxis,:] + arange(ntr)[:,newaxis])
lcids = repeat(arange(ntr), npt)
epids = arange(ntr)

fig, axs = subplots(3, 1, figsize=(13,6), sharex='all')
axs[0].plot(times)
axs[1].plot(lcids)
axs[2].plot(epids[lcids])
setp(axs[0], ylabel='Time [d]')
setp(axs[1], ylabel='Light curve index')
setp(axs[2], ylabel='Epoch index', xlabel='Point')
fig.tight_layout()

"""
Explanation: Set up the model
First, we create the time, light curve index, and epoch index arrays. We model ntr consecutive transits, each with npt points.
End of explanation
"""


tm.set_data(times, lcids=lcids)

flux = tm.evaluate(k=0.11, ldc=[0.23, 0.1], t0=1.0, p=1.0, a=4.0, i=0.5*pi)

fig, ax = subplots()
ax.plot(times, flux, 'k.-')
fig.tight_layout()

"""
Explanation: Model evaluation for a linear ephemeris
First, we check that everything works as it should when we assume that the orbit does not evolve over time.
End of explanation
"""


tm.set_data(times, lcids=lcids, epids=epids)

flux = tm.evaluate(k=0.11, ldc=[0.23, 0.1], t0=[0.0, 1.1, 1.9], p=1.0, a=4.0, i=0.5*pi)

fig, ax = subplots()
ax.plot(times, flux, 'k.-')
fig.tight_layout()

"""
Explanation: TTV model evaluation
Next, we give the TransitModel.set_data method the per-light-curve epoch indices (epids), and then a separate transit centre time in the model evaluation, one per light curve.
End of explanation
"""


npv = 10
ks = normal(0.1, 0.01, npv)
ldcs = tile([0.23, 0.1], [npv,1,1])
t0s = tile([0.0, 1.1, 1.9], [npv, 1])
t0s[:,1] += normal(0.0, 0.02, npv)
ps = full(npv, 1.0)
smas = full(npv, 4.0)
incs = full(npv, 0.5*pi)

flux = tm.evaluate(k=ks, ldc=ldcs, t0=t0s, p=ps, a=smas, i=incs)

fig, ax = subplots()
ax.plot(times, flux.T, '-')
fig.tight_layout()

"""
Explanation: As with all the PyTransit models, the TTV model can be evaluated for a set of parameter vectors in a single call.
End of explanation
"""


