tricycle/lightcurve.py at master · StellarArmy/tricycle · GitHub

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
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
from gatspy.periodic import LombScargleFast
import matplotlib.pyplot as plt

class LightCurve(object):
    """
    Eclipsing binary light curve with system parameters.

    Parent class with analysis methods.

    Parameters
    ----------
    time : ndarray
        The observation times, in days.
    flux : ndarray
        The system flux. May be in physical or relative units.
    err : ndarray
        The flux errors. Should have same units as `flux`.
    p_orb : float
        Orbital period in days.
    t_0 : float
        Reference time of mid-eclipse in days.
    p_depth : float
        Depth of primary eclipse.
    s_depth : float
        Depth of secondary eclipse.
    p_width : float
        Width of primary eclipse in phase.
    s_width : float
        Width of secondary eclipse in phase.
    sep : float
        Seperation between the primary eclipse in secondary eclipse in phase
    kic : int
        Kepler Input Catalog (KIC) ID number.

    """
    def __init__(self, time, flux, err, p_orb, t_0, p_depth, s_depth, p_width,
                 s_width, sep, kic):

        self.time = time
        self.flux = flux
        self.err = err
        self.p_orb = p_orb
        self.t_0 = t_0
        self.p_depth = p_depth
        self.s_depth = s_depth
        self.p_width = p_width
        self.s_width = s_width
        self.sep = sep
        self.kic = kic

    def get_period(self, units='days'):
        """
        Return the orbital period.

        Parameters
        ----------
        units : {'days', 'hours', 'minutes'}
            Desired units for period.

        Returns
        -------
        period : float
            The orbital period in the specified units.

        """
        if units == 'days':
            period = self.p_orb
        elif units == 'hours':
            period = self.p_orb * 24.
        elif units == 'minutes':
            period = self.p_orb * 24. * 60.
        else:
            raise ValueError('Invalid choice of units.')

        return period

    def curve_cut(self):
        """This a function that takes parameters of a star and gives you the time and flux with the eclipses cut out

        Parameters:
        Time
        Flux
        bjdo
        period
        pwidth
        swidth
        separation

        return:
            timecut and fluxcut"""
        # Convert BJD to Kepler date
        t_0 = self.t_0 - 54833

        phase=((self.time- t_0) % self.p_orb) / self.p_orb
        mask= ((phase > self.p_width) & (phase < 1- self.p_width)) & ((phase > self.sep+ self.s_width) | (phase < self.sep- self.s_width))

        timecut= self.time[mask]
        fluxcut= self.flux[mask]
        errcut = self.err[mask]

        return timecut, fluxcut, errcut

    def periodogram(self, time, flux, err, p_fold=None, plt_color='k',
                    max_days=100.0, oversampling=5):
        """
        Plot a the light curve, periodogram, and phase-folded light curve.

        The orbital period and possible aliases are indicated in red on the
        periodogram.

        Parameters
        ----------
        time : array_like
            Observation times in days.
        flux : array_like
            Fluxes.
        err : array_like
            Flux errors.
        p_fold : float, optional
            Plot the light curve folded at this period (in days), and indicate
            the period and likely aliases on the periodogram plot.
        plt_color : str, optional
            The line and scatter plot color.
        max_days : float, optional
            The maximum number of days to plot in the light curve and
            periodogram.
        oversampling: int, optional
            The oversampling factor for the periodogram.

        """
        model = LombScargleFast().fit(time, flux, err)
        period, power = model.periodogram_auto(oversampling=oversampling)

        fig1, (ax1, ax2) = plt.subplots(nrows=2, figsize=(7, 14))

        ax1.plot(time, flux, color=plt_color)
        ax1.set_xlim(time.min(), time.min() + max_days)
        ax1.set_xlabel('Time (days)')
        ax1.set_ylabel('Relative Flux')

        ax2.plot(period, power, color=plt_color)
        ax2.set_xlim(0.1, max_days)
        ax2.set_xlabel('Period (days)')
        ax2.set_ylabel('Power')

        # Plot some the most likely aliases of eclipse period.
        factors = [0.5, 1, 2, 3, 4]
        for factor in factors:
            ax2.axvline(self.p_orb / factor, color='r')

        if p_fold is not None:
            for factor in factors:
                ax2.axvline(p_fold / factor, color='b')

        fig1.suptitle('KIC {0:d} --- p_orb = {1:3.5f} days'.
                      format(self.kic, self.p_orb))

        if p_fold is not None:
            fig2, ax3 = plt.subplots()
            phase = (time % p_fold) / p_fold
            ax3.scatter(phase, flux, color=plt_color, s=0.1)
            ax3.set_xlim(0, 1)
            ax3.set_xlabel('Phase')
            ax3.set_ylabel('Relative Flux')

        plt.show()