%matplotlib inline
import numpy as np
import matplotlib.pyplot as plt

# use seaborn's default plotting styles for matplotlib
import seaborn; seaborn.set()

from gatspy.datasets import fetch_rrlyrae
rrlyrae = fetch_rrlyrae()
lcid = rrlyrae.ids[0]

t, y, dy, filts = rrlyrae.get_lightcurve(lcid)
period = rrlyrae.get_metadata(lcid)['P']

for filt in 'ugriz':
    mask = (filts == filt)
    plt.errorbar(t[mask] % period, y[mask], dy[mask], fmt='.', label=filt)
plt.gca().invert_yaxis()
plt.legend(ncol=3, loc='upper left');

"""
Explanation: Multiband Periodograms with gatspy
gatspy includes an implementation of the multiband periodogram, introduced in VanderPlas & Ivezic 2015. If you've taken a look at the Single Band Example, the interface should be pretty familiar.
As in that example, let's start with some imports, and downloading some of the Sesear 2010 Stripe 82 RR Lyrae data, and plotting it:
End of explanation
"""


from gatspy.periodic import LombScargleMultiband
model = LombScargleMultiband(Nterms_base=1, Nterms_band=0)
model.fit(t, y, dy, filts)
periods = np.linspace(period - 0.1, period + 0.1, 2000)
power = model.periodogram(periods)
plt.plot(periods, power, lw=1)
plt.xlim(periods[0], periods[-1]);

"""
Explanation: Now we'll fit the multiband periodogram model to this data. For more information on the model, refer to the VanderPlas and Ivezic paper mentioned above.
End of explanation
"""


def plot_model(model, lcid):
    t, y, dy, filts = rrlyrae.get_lightcurve(lcid)
    model.fit(t, y, dy, filts)
    
    tfit = np.linspace(0, period, 1000)
    for filt in 'ugriz':
        mask = (filts == filt)
        eb = plt.errorbar(t[mask] % period, y[mask], dy[mask], fmt='.', label=filt)
        yfit = model.predict(tfit, filt, period=period)
        plt.plot(tfit, yfit, color=eb[0].get_color())
    plt.gca().invert_yaxis()
    plt.legend(ncol=3, loc='upper left')
    
plot_model(LombScargleMultiband(Nterms_base=1, Nterms_band=0), lcid)

"""
Explanation: We can see what the multiterm model looks like by plotting it over the data:
End of explanation
"""


plot_model(LombScargleMultiband(Nterms_base=4, Nterms_band=1), lcid)

"""
Explanation: If we'd like to do a higher-oder multiterm model, we can simply adjust the number of terms in the base and band models:
End of explanation
"""


