def calcBls(flux, time, bls_durs, minP=None, maxP=None, min_trans=3):
"""
Take a bls and return the spectrum.
"""
bls = BoxLeastSquares(time, flux)
period_grid = bls.autoperiod(bls_durs,minimum_period=minP, \
maximum_period=maxP, minimum_n_transit=min_trans, \
frequency_factor=0.8)
bls_power = bls.power(period_grid, bls_durs, oversample=20)
return bls_power
def find_and_mask_transits(time,
flux,
flux_err,
periods,
durations,
nplanets=1,
plot=False):
"""
Iteratively find and mask transits in the flattened light curve.
Args:
time (array): The time array.
flux (array): The flux array. You'll get the best results
if this is flattened.
flux_err (array): The array of flux uncertainties.
periods (array): The array of periods to search over for BLS.
For example, periods = np.linspace(0.5, 20, 10)
durations (array): The array of durations to search over for BLS.
For example, durations = np.linspace(0.05, 0.2, 10)
nplanets (Optional[int]): The number of planets you'd like to search for.
This function will interatively find and remove nplanets. Default is 1.
Returns:
transit_masks (list): a list of masks that correspond to the in
transit points of each light curve. To mask out transits do
time[~transit_masks[index]], etc.
"""
cum_transit = np.ones(len(time), dtype=bool)
_time, _flux, _flux_err = time * 1, flux * 1, flux_err * 1
t0s, durs, porbs = [np.zeros(nplanets) for i in range(3)]
transit_masks = []
for i in range(nplanets):
bls = BoxLeastSquares(t=_time, y=_flux, dy=_flux_err)
bls.power(periods, durations)
periods = bls.autoperiod(durations,
minimum_n_transit=3,
frequency_factor=5.0)
results = bls.autopower(durations, frequency_factor=5.0)
# Find the period of the peak
period = results.period[np.argmax(results.power)]
# Extract the parameters of the best-fit model
index = np.argmax(results.power)
porbs[i] = results.period[index]
t0s[i] = results.transit_time[index]
durs[i] = results.duration[index]
if plot:
# Plot the periodogram
fig, ax = plt.subplots(1, 1, figsize=(10, 5))
ax.plot(results.period, results.power, "k", lw=0.5)
ax.set_xlim(results.period.min(), results.period.max())
ax.set_xlabel("period [days]")
ax.set_ylabel("log likelihood")
# Highlight the harmonics of the peak period
ax.axvline(period, alpha=0.4, lw=4)
for n in range(2, 10):
ax.axvline(n * period, alpha=0.4, lw=1, linestyle="dashed")
ax.axvline(period / n, alpha=0.4, lw=1, linestyle="dashed")
# plt.show()
# plt.plot(_time, _flux, ".")
# plt.xlim(1355, 1360)
in_transit = bls.transit_mask(_time, porbs[i], durs[i], t0s[i])
transit_masks.append(in_transit)
_time, _flux, _flux_err = _time[~in_transit], _flux[~in_transit], \
_flux_err[~in_transit]
return transit_masks, t0s, durs, porbs
def from_lightcurve(lc, **kwargs):
"""Creates a Periodogram from a LightCurve using the Box Least Squares (BLS) method."""
# BoxLeastSquares was added to `astropy.stats` in AstroPy v3.1 and then
# moved to `astropy.timeseries` in v3.2, which makes the import below
# somewhat complicated.
try:
from astropy.timeseries import BoxLeastSquares
except ImportError:
try:
from astropy.stats import BoxLeastSquares
except ImportError:
raise ImportError("BLS requires AstroPy v3.1 or later")
# Validate user input for `lc`
# (BoxLeastSquares will not work if flux or flux_err contain NaNs)
lc = lc.remove_nans()
if np.isfinite(lc.flux_err).all():
dy = lc.flux_err
else:
dy = None
# Validate user input for `duration`
duration = kwargs.pop("duration", 0.25)
if duration is not None and ~np.all(np.isfinite(duration)):
raise ValueError("`duration` parameter contains illegal nan or inf value(s)")
# Validate user input for `period`
period = kwargs.pop("period", None)
minimum_period = kwargs.pop("minimum_period", None)
maximum_period = kwargs.pop("maximum_period", None)
if period is not None and ~np.all(np.isfinite(period)):
raise ValueError("`period` parameter contains illegal nan or inf value(s)")
if minimum_period is None:
if period is None:
minimum_period = np.max([np.median(np.diff(lc.time)) * 4,
np.max(duration) + np.median(np.diff(lc.time))])
else:
minimum_period = np.min(period)
if maximum_period is None:
if period is None:
maximum_period = (np.max(lc.time) - np.min(lc.time)) / 3.
else:
maximum_period = np.max(period)
# Validate user input for `time_unit`
time_unit = (kwargs.pop("time_unit", "day"))
if time_unit not in dir(u):
raise ValueError('{} is not a valid value for `time_unit`'.format(time_unit))
# Validate user input for `frequency_factor`
frequency_factor = kwargs.pop("frequency_factor", 10)
df = frequency_factor * np.min(duration) / (np.max(lc.time) - np.min(lc.time))**2
npoints = int(((1/minimum_period) - (1/maximum_period))/df)
if npoints > 1e7:
raise ValueError('`period` contains {} points.'
'Periodogram is too large to evaluate. '
'Consider setting `frequency_factor` to a higher value.'
''.format(np.round(npoints, 4)))
elif npoints > 1e5:
log.warning('`period` contains {} points.'
'Periodogram is likely to be large, and slow to evaluate. '
'Consider setting `frequency_factor` to a higher value.'
''.format(np.round(npoints, 4)))
# Create BLS object and run the BLS search
bls = BoxLeastSquares(lc.time, lc.flux, dy)
if period is None:
period = bls.autoperiod(duration,
minimum_period=minimum_period,
maximum_period=maximum_period,
frequency_factor=frequency_factor)
result = bls.power(period, duration, **kwargs)
if not isinstance(result.period, u.quantity.Quantity):
result.period = u.Quantity(result.period, time_unit)
if not isinstance(result.power, u.quantity.Quantity):
result.power = result.power * u.dimensionless_unscaled
if not isinstance(result.duration, u.quantity.Quantity):
result.duration = u.Quantity(result.duration, time_unit)
return BoxLeastSquaresPeriodogram(frequency=1. / result.period,
power=result.power,
default_view='period',
label=lc.label,
targetid=lc.targetid,
transit_time=result.transit_time,
duration=result.duration,
depth=result.depth,
bls_result=result,
snr=result.depth_snr,
bls_obj=bls,
time=lc.time,
flux=lc.flux,
time_unit=time_unit)
def bls_estimator(
x,
y,
yerr=None,
duration=0.2,
min_period=None,
max_period=None,
objective=None,
method=None,
oversample=10,
**kwargs,
):
"""Estimate the period of a time series using box least squares
All extra keyword arguments are passed directly to
:func:`astropy.timeseries.BoxLeastSquares.autopower`.
Args:
x (ndarray[N]): The times of the observations
y (ndarray[N]): The observations at times ``x``
yerr (Optional[ndarray[N]]): The uncertainties on ``y``
min_period (Optional[float]): The minimum period to consider
max_period (Optional[float]): The maximum period to consider
Returns:
A dictionary with the computed autocorrelation function and the
estimated period. For compatibility with the
:func:`lomb_scargle_estimator`, the period is returned as a list with
the key ``peaks``.
"""
kwargs["minimum_period"] = kwargs.get("minimim_period", min_period)
kwargs["maximum_period"] = kwargs.get("maximum_period", max_period)
x_ref = 0.5 * (np.min(x) + np.max(x))
bls = BoxLeastSquares(x - x_ref, y, yerr)
# Estimate the frequency factor to not be insanely slow
if "frequency_factor" not in kwargs:
kwargs["frequency_factor"] = 1.0
periods = bls.autoperiod(duration, **kwargs)
while len(periods) > len(x):
kwargs["frequency_factor"] *= 2
periods = bls.autoperiod(duration, **kwargs)
# Compute the periodogram
pg = bls.autopower(
duration,
objective=objective,
method=method,
oversample=oversample,
**kwargs,
)
# Correct for the reference time offset
pg.transit_time += x_ref
# Find the peak
peaks = find_peaks(1 / pg.period, pg.power, max_peaks=1)
results = dict(bls=pg, peaks=peaks, peak_info=None)
if not len(peaks):
return results
# Extract the relevant information at the peak
ind = peaks[0]["index"]
results["peak_info"] = dict(
(k, v[ind]) for k, v in pg.items() if k != "objective")
return results
