def fit_sine(time, flux, unc, period):
"""Fit a simple sine model fixed to the best-fit period."""
def _sine_model(t, amp, yoffset, tshift):
return amp * np.sin(2 * np.pi * t / period + tshift) + yoffset
# Phase by the period, then extend the arrays
# To fit two cycles instead of one
phased_time = utils.phase(time, period)
fit_time = np.append(phased_time, phased_time + period)
fit_flux = np.append(flux, flux)
fit_unc = np.append(unc, unc)
# initial amplitude and yoffset are stdev and median, respectively
p0 = np.append(utils.stats(flux, unc), 0.0)
popt, pcov = opt.curve_fit(_sine_model,
fit_time,
fit_flux,
sigma=fit_unc,
p0=p0)
perr = np.sqrt(np.diag(pcov))
logging.debug("amplitude, yoffset, tshift")
logging.debug(popt)
logging.debug(perr)
return phased_time, _sine_model(phased_time, *popt)
def fit_sine(time, flux, unc, period):
"""Fit a simple sine model fixed to the best-fit period."""
def _sine_model(t, amp, yoffset, tshift):
return amp * np.sin(2 * np.pi * t / period + tshift) + yoffset
# Phase by the period, then extend the arrays
# To fit two cycles instead of one
phased_time = utils.phase(time, period)
fit_time = np.append(phased_time, phased_time + period)
fit_flux = np.append(flux, flux)
fit_unc = np.append(unc, unc)
# initial amplitude and yoffset are stdev and median, respectively
p0 = np.append(utils.stats(flux, unc), 0.0)
popt, pcov = opt.curve_fit(_sine_model, fit_time, fit_flux, sigma=fit_unc, p0=p0)
perr = np.sqrt(np.diag(pcov))
logging.debug("amplitude, yoffset, tshift")
logging.debug(popt)
logging.debug(perr)
return phased_time, _sine_model(phased_time, *popt)
def pre_whiten(time, flux, unc_flux, period, kind="supersmoother",
which="phased",phaser=None):
"""Phase a lightcurve and then smooth it.
Inputs
------
time, flux, unc_flux: array_like
period: float
period to whiten by
kind: string, optional
type of smoothing to use. Defaults to "supersmoother."
Other types are "boxcar", "linear"
which: string, optional
whether to smooth the "phased" lightcurve (default) or the "full"
lightcurve.
phaser: Float, optional (default=None)
if kind="boxcar", phaser is the Half-width of the smoothing window.
if kind="supersmoother", phaser is alpha (the "bass enhancement").
Outputs
-------
white_flux, white_unc, smoothed_flux: arrays
"""
# phase the LC by the period
if period is not None:
phased_time = utils.phase(time, period)
else:
phased_time = time
if kind.lower()=="supersmoother":
if phaser is None:
logging.info("Phaser not set! "
"Set phaser=alpha (bass-enhancement value "
"for supersmoother) if desired.")
if which.lower()=="phased":
# Instantiate the supersmoother model object with the input period
model = supersmoother.SuperSmoother(period=period)
# Set up base arrays for phase for the fit
x_vals = np.linspace(0,max(phased_time),1000)
# Run a fit for the y phase
y_fit = model.fit(phased_time, flux, unc_flux).predict(x_vals)
elif which.lower()=="full":
# Instantiate the supersmoother model object with the input period
model = supersmoother.SuperSmoother(alpha=phaser)
# Set up base arrays for time for the fit
x_vals = np.linspace(min(time),max(time),1000)
# run a fit for the y values
y_fit = model.fit(time, flux, unc_flux).predict(x_vals)
else:
logging.warning("unknown which %s !!!",which)
elif kind.lower()=="boxcar":
if phaser is None:
logging.info("Phaser not set! "
"Set phaser to the width of the smoothing "
"box in pixels!")
if which.lower()=="phased":
# sort the phases
sort_locs = np.argsort(phased_time)
x_vals = phased_time[sort_locs]
flux_to_fit = flux[sort_locs]
elif which.lower()=="full":
x_vals = time
flux_to_fit = flux
else:
logging.warning("unknown which %s !!!",which)
# Use astropy's convolution function!
boxcar_kernel = convolution.Box1DKernel(width=phaser,
mode="linear_interp")
y_fit = convolution.convolve(flux_to_fit, boxcar_kernel,
boundary="wrap")
elif kind=="linear":
if which!="full":
logging.warning("Linear smoothing only allowed for full "
"lightcurve! Switching to full mode.")
which = "full"
# Fit a line to the data
pars = np.polyfit(time, flux, deg=1)
m, b = pars
smoothed_flux = m * time + b
else:
logging.warning("unknown kind %s !!!",kind)
if (kind=="supersmoother") or (kind=="boxcar"):
# Interpolate back onto the original time array
interp_func = interpolate.interp1d(x_vals, y_fit)
if which.lower()=="phased":
smoothed_flux = interp_func(phased_time)
elif which.lower()=="full":
smoothed_flux = interp_func(time)
else:
logging.warning("unknown which %s !!!",which)
# Whiten the input flux by subtracting the smoothed version
# The smoothed flux represents the "bulk trend"
white_flux = flux - smoothed_flux
white_unc = unc_flux
return white_flux, white_unc, smoothed_flux
def plot_one(lightcurve,
periodogram,
best_period,
power_threshold,
data_label,
residuals=None,
aliases=None,
axes_list=None,
plot_title=None,
phase_by=None,
**plot_kwargs):
"""Plot input lightcurve, input periodogram, and input phased lightcurve
INPUTS ARE DIFFERENT FROM REST OF CODE
inputs
------
lightcurve: arraylike (3, num_points)
time, flux, unc_flux
periodogram: arraylike (2, num_periods)
periods_to_test, periodogram
best_period: float
power_threshold: float
data_label: string
name for the dataset being plotted
residuals: array-like (default None)
if provided, will be plotted in the bottom panel
aliases: array-like (default None)
if provided, will be plotted as vertical dotted lines on the
periodogram
axes_list: optional, length=4
if not provided, a new figure will be created.
Axes should correspond to lightcurve, periodogram, phased lightcurve,
and phased residuals
phase_by: float (default=None)
if not None, the phased lightcurve will use this period
instead of best_period
"""
# If axes_list is None, create the figure and axes with setup_plots
if axes_list is None:
fig, axes_list = setup_plots()
plot_color, plot_marker = color1, shape1
elif len(axes_list) != 4:
logging.warning("Incorrect number of axes! Setting up new figure")
fig, axes_list = setup_plots()
plot_color, plot_marker = color1, shape1
else:
fig = plt.gcf()
plot_color, plot_marker = color2, shape2
# Top panel: lightcurve
axes_list[0].plot(lightcurve[0],
lightcurve[1],
lw=0,
marker=plot_marker,
mec=plot_color,
mfc=plot_color,
ms=2,
label=data_label)
# Also plot vertical lines for the period if it's longer than 2 days
# (shorter than that isn't really comprehensible)
if best_period >= 2.0:
phase_mult = np.arange(min(lightcurve[0]), max(lightcurve[0]),
best_period)
for pm in phase_mult:
axes_list[0].axvline(pm,
color=plot_color,
linestyle="--",
zorder=-100,
alpha=0.75)
# Middle panel: periodogram
logging.debug("plot periodograms")
axes_list[1].plot(periodogram[0], periodogram[1], color=plot_color)
axes_list[1].axvline(best_period, color=plot_color, linestyle="--")
axes_list[1].set_xlim(xmin=0.1)
# Plot the power threshold, if it would be visible
ymax = axes_list[1].get_ylim()[1]
logging.debug("ymax %f", ymax)
if (((type(power_threshold) == float) or (type(power_threshold) == int))
and power_threshold < ymax):
# Only one threshold
axes_list[1].axhline(power_threshold, color="Grey", ls="-.")
elif (type(power_threshold) == float and power_threshold >= ymax):
logging.debug("peak lower than threshold")
else:
for threshold in power_threshold:
logging.debug(threshold)
axes_list[1].axhline(threshold, color=plot_color, ls="-.")
# Plot aliases, if provided
if aliases is not None:
for alias in aliases:
axes_list[1].axvline(alias, color=plot_color, linestyle=":")
# Plot harmonics of the best period
harmonics = np.array([0.5, 2]) * best_period
for harm in harmonics:
axes_list[1].axvline(harm, color=plot_color, linestyle="--", alpha=0.5)
# Plot harmonics of the thruster firing time
harmonics = np.append(0.125, np.arange(0.25, 2.1, 0.25))
for harm in harmonics:
axes_list[1].axvline(harm,
color="LightGrey",
linestyle="-",
alpha=0.5,
zorder=-111)
# Bottom panels: phased lightcurve
if phase_by is None:
phase_by = best_period
phased = utils.phase(lightcurve[0], phase_by)
axes_list[2].plot(phased,
lightcurve[1],
lw=0,
marker=plot_marker,
mec=plot_color,
mfc=plot_color,
ms=2)
axes_list[2].set_xlim(0, phase_by)
if residuals is not None:
axes_list[3].plot(phased,
residuals,
lw=0,
marker=plot_marker,
mec=plot_color,
mfc=plot_color,
ms=2)
axes_list[3].set_xlim(0, phase_by)
# plt.tight_layout()
return fig, axes_list
def plot_one(lightcurve, periodogram, best_period, power_threshold, data_label,
residuals=None, aliases=None, axes_list=None,plot_title=None,
phase_by=None, **plot_kwargs):
"""Plot input lightcurve, input periodogram, and input phased lightcurve
INPUTS ARE DIFFERENT FROM REST OF CODE
inputs
------
lightcurve: arraylike (3, num_points)
time, flux, unc_flux
periodogram: arraylike (2, num_periods)
periods_to_test, periodogram
best_period: float
power_threshold: float
data_label: string
name for the dataset being plotted
residuals: array-like (default None)
if provided, will be plotted in the bottom panel
aliases: array-like (default None)
if provided, will be plotted as vertical dotted lines on the
periodogram
axes_list: optional, length=4
if not provided, a new figure will be created.
Axes should correspond to lightcurve, periodogram, phased lightcurve,
and phased residuals
phase_by: float (default=None)
if not None, the phased lightcurve will use this period
instead of best_period
"""
# If axes_list is None, create the figure and axes with setup_plots
if axes_list is None:
fig, axes_list = setup_plots()
plot_color, plot_marker = color1, shape1
elif len(axes_list)!=4:
logging.warning("Incorrect number of axes! Setting up new figure")
fig, axes_list = setup_plots()
plot_color, plot_marker = color1, shape1
else:
fig = plt.gcf()
plot_color, plot_marker = color2, shape2
# Top panel: lightcurve
axes_list[0].plot(lightcurve[0], lightcurve[1], lw=0,
marker=plot_marker, mec=plot_color,
mfc=plot_color, ms=2, label=data_label)
# Also plot vertical lines for the period if it's longer than 2 days
# (shorter than that isn't really comprehensible)
if best_period>=2.0:
phase_mult = np.arange(min(lightcurve[0]), max(lightcurve[0]),
best_period)
for pm in phase_mult:
axes_list[0].axvline(pm, color=plot_color, linestyle="--",
zorder=-100, alpha=0.75)
# Middle panel: periodogram
logging.debug("plot periodograms")
axes_list[1].plot(periodogram[0], periodogram[1], color=plot_color)
axes_list[1].axvline(best_period, color=plot_color, linestyle="--")
axes_list[1].set_xlim(xmin=0.1)
# Plot the power threshold, if it would be visible
ymax = axes_list[1].get_ylim()[1]
logging.debug("ymax %f", ymax)
if (((type(power_threshold)==float) or (type(power_threshold)==int))
and power_threshold<ymax):
# Only one threshold
axes_list[1].axhline(power_threshold,color="Grey",ls="-.")
elif (type(power_threshold)==float and
power_threshold>=ymax):
logging.debug("peak lower than threshold")
else:
for threshold in power_threshold:
logging.debug(threshold)
axes_list[1].axhline(threshold,color=plot_color,ls="-.")
# Plot aliases, if provided
if aliases is not None:
for alias in aliases:
axes_list[1].axvline(alias, color=plot_color, linestyle=":")
# Plot harmonics of the best period
harmonics = np.array([0.5, 2])*best_period
for harm in harmonics:
axes_list[1].axvline(harm, color=plot_color, linestyle="--",
alpha=0.5)
# Plot harmonics of the thruster firing time
harmonics = np.append(0.125,np.arange(0.25,2.1,0.25))
for harm in harmonics:
axes_list[1].axvline(harm, color="LightGrey", linestyle="-",
alpha=0.5, zorder=-111)
# Bottom panels: phased lightcurve
if phase_by is None:
phase_by = best_period
phased = utils.phase(lightcurve[0], phase_by)
axes_list[2].plot(phased, lightcurve[1], lw=0, marker=plot_marker,
mec=plot_color, mfc=plot_color, ms=2)
axes_list[2].set_xlim(0,phase_by)
if residuals is not None:
axes_list[3].plot(phased, residuals, lw=0, marker=plot_marker,
mec=plot_color, mfc=plot_color, ms=2)
axes_list[3].set_xlim(0,phase_by)
# plt.tight_layout()
return fig, axes_list
