/
lightcurve.py
executable file
·106 lines (100 loc) · 4.18 KB
/
lightcurve.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import numpy
import orbit
import ma02
import scipy.optimize
from matplotlib import pyplot as plt
def lightcurve(JD, P, p, Ttr = 0, Ecc = 0, a = 10, \
incl = numpy.pi/2, \
omega = 0, limbd = [0, 0], sec = 0):
"""Calculate transit / eclipse light curve according to Mandel &
Agol (2002), with linear, quadratic or 4-parameter non-linear limb-darkening.
TSB: modified the code to calculate the model on a finer grid than the
original data is sampled on. The model is then averaged over the finer data.
This accounts for non-linear changes in the brightnesses during ingress
and egress. """
#try binning model
JD = numpy.array(JD)
avediff = numpy.median(JD[1:] - JD[0:-1])
#add in 4 extra data points per obs
JDnew = numpy.zeros(len(JD)*5)
flagorig = numpy.zeros(len(JD)*5)
for i in numpy.arange(len(JD)):
#i*5 + 2
JDnew[(i*5)] = JD[i] - (2.*avediff*0.25)
JDnew[(i*5)+1] = JD[i] - (avediff*0.25)
JDnew[(i*5)+2] = JD[i]
JDnew[(i*5)+3] = JD[i] + (avediff*0.25)
JDnew[(i*5)+4] = JD[i] + (2.*avediff*0.25)
flagorig[(i*5)+2] = 1.
JDorig = numpy.copy(JD)
JD = JDnew
T0 = orbit.getT0(P, Ttr, omega, Ecc)
x,y,z = orbit.skycoord(JD, P, T0, Ecc, a, incl, 0, omega)
r = numpy.sqrt(x**2+y**2)
res = numpy.ones(len(JD))
f_uni = 1.0 - ma02.occultuniform(r, p)
if numpy.size(limbd) == 0:
f_tr = f_uni
elif numpy.size(limbd) == 4:
f_tr = ma02.occultnonlin(r, p, limbd)
else:
f_tr = ma02.occultquad(r, p, limbd)
f_sec = 1.0 - (1.0 - f_uni) * sec / p**2
#now sum up model
f_trnew = numpy.zeros(len(JDorig))
f_secnew = numpy.zeros(len(JDorig))
znew = numpy.zeros(len(JDorig))
for i in numpy.arange(len(JDorig)):
f_trnew[i] = (f_tr[(i*5)] + f_tr[(i*5)+1] + f_tr[(i*5)+2] +
f_tr[(i*5)+3] + f_tr[(i*5)+4]) / 5.
f_secnew[i] = (f_sec[(i*5)] + f_sec[(i*5)+1] + f_sec[(i*5)+2] +
f_sec[(i*5)+3] + f_sec[(i*5)+4]) / 5.
znew[i] = z[(i*5)+2]
f_tr = f_trnew
f_sec = f_secnew
z = znew
return numpy.select([z > 0, z <= 0], [f_tr, f_sec])
def transit_errfunc_ptia(par_tofit, par_fixed, JD, flux, flux_err):
p, Ttr, incl, a = par_tofit
P, Ecc, omega, limbd, sec = par_fixed
flux_model = lightcurve(JD, P, p, Ttr, Ecc, a, incl, omega, limbd, sec)
chi2 = numpy.sum((flux - flux_model)**2 / flux_err**2)
print 'chi2: ', chi2
return chi2
def transit_fit_example():
'''Simulate some data containing a transit plus white Gaussian
noise and fit for the planet / star radius ratio (p = Rp / Rs),
epoch Ttr, inclination, and ratio of semi-major axis to stellar
radius (a/Rs), while holding the other parameters fixed at their
true values.'''
JD = numpy.arange(100) * 0.01 - 0.5 + numpy.random.normal(0, 0.005, 100)
P = 3.0
Ttr_true = 0.0; Ttr_guess = 0.02
p_true = 0.1; p_guess = 0.08
incl_true = 0.99 * numpy.pi / 2.; incl_guess = numpy.pi / 2.
a_true = 5.0; a_guess = 4.5
cn = [0, 0.45, 0, 0.05]
Ecc = 0; omega = 0; sec = 0
truth = lightcurve(JD, P, p_true, Ttr_true, Ecc, a_true, incl_true, \
omega, cn, sec)
flux = truth + nupmy.random.normal(0, 0.002, numpy.size(JD))
flux_err = numpy.ones(numpy.size(JD)) * 0.002
initial_guess = lightcurve(JD, P, p_guess, Ttr_guess, \
Ecc, a_guess, incl_guess, omega, cn, sec)
p_fit, Ttr_fit, incl_fit, a_fit = \
scipy.optimize.fmin(transit_errfunc_ptia, \
[p_guess, Ttr_guess, incl_guess, a_guess], \
args = ([P, Ecc, omega, cn, sec], JD, flux, flux_err))
fit = lightcurve(JD, P, p_fit, Ttr_fit, \
Ecc, a_fit, incl_fit, omega, cn, sec)
plt.clf()
plt.errorbar(JD, flux, yerr = flux_err, fmt = 'k.', label = 'data')
plt.plot(JD, truth, 'k-', label = 'true')
plt.plot(JD, initial_guess, 'b--', label = 'initial')
plt.plot(JD, fit, 'r-', label = 'fit')
plt.legend(loc = 'lower right')
print p_true, p_guess, p_fit
print Ttr_true, Ttr_guess, Ttr_fit
print incl_true, incl_guess, incl_fit
print a_true, a_guess, a_fit
return