def test_plot_image():
mu = [3.5, 3]
x, y = np.mgrid[0:15, 0:15]
pos = np.dstack((x, y))
var = multivariate_normal(mean=mu, cov=[[1, 0], [0, 1]])
gauss = 5 * var.pdf(
pos) - 0.05 # Make sure it has some negative values as well
gauss[8, 8] = np.NaN
gauss[4, 4] = -0.2
scales = ['linear', 'sqrt', 'log', 'asinh', 'histeq', 'sinh', 'squared']
fig, axes = plt.subplots(2, 4, figsize=(14, 8))
axes = axes.flatten()
for k, scale in enumerate(scales):
ax = axes[k]
plot_image(gauss, ax=ax, scale=scale, title=scale, cbar='right')
ax.plot(mu[1], mu[0], 'r+')
# In the final plot:
plot_image(gauss,
ax=axes[-1],
scale='log',
title='log - Reds',
cmap='Reds',
cbar='right')
fig.tight_layout()
return fig
def test_plot_cbar_and_nans():
# Construct image:
np.random.seed(42)
img = np.random.rand(10, 10)
img[2:8, 2:8] = np.NaN
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, figsize=(16, 6))
plot_image(img, ax=ax1, scale='linear', vmin=0.4, cbar='left')
plot_image(img, ax=ax2, scale='sqrt', vmin=0.4, cbar='bottom')
plot_image(img, ax=ax3, scale='log', vmin=0.4, cbar='right')
plot_image(img, ax=ax4, scale='asinh', vmin=0.4, cbar='top')
return fig
def test_cbv_fit(INPUT_DIR):
# Create CBV object:
cbv = CBV(
os.path.join(INPUT_DIR, 'cbv-prepare', 'cbv-s0001-c1800-a143.hdf5'))
coeffs = [10, 500, 50, 100, 0, 10, 0]
abs_flux = 3500
# Create model using coefficients, and make fake lightcurve out of it:
mdl = cbv.mdl(coeffs) * abs_flux
# Another check of making crazy weights:
#sigma = np.ones_like(mdl)*100
#mdl[200] = 50000
#sigma[200] = 1e-17
lc = TessLightCurve(time=cbv.time, flux=mdl)
# Run CBV fitting with fixed number of CBVs:
flux_filter, res, diagnostics = cbv.fit(lc, cbvs=3, use_bic=False)
# Plot:
fig, (ax1, ax2) = plt.subplots(2, 1, figsize=(12, 16))
ax1.scatter(cbv.time, mdl, alpha=0.3)
ax1.plot(cbv.time, flux_filter, alpha=0.5, color='r')
ax2.scatter(cbv.time, mdl - flux_filter)
# Check the diagnostics dict:
print(diagnostics)
assert diagnostics['method'] == 'LS'
assert not diagnostics['use_bic']
assert not diagnostics['use_prior']
# Check the coefficients coming out of the fit:
# They should be the same as the ones we put in
print(res - coeffs)
np.testing.assert_allclose(res, coeffs, atol=0.5, rtol=0.5)
# The fitted model should be very close to the model going in:
np.testing.assert_allclose(mdl, flux_filter)
def test_plot_image_invalid():
mu = [3.5, 3]
x, y = np.mgrid[0:10, 0:10]
pos = np.dstack((x, y))
var = multivariate_normal(mean=mu, cov=[[1, 0], [0, 1]])
gauss = var.pdf(pos)
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 6))
# Run with invalid scale:
with pytest.raises(ValueError):
plot_image(gauss, ax=ax1, scale='invalid-scale')
# Plot with single NaN:
gauss[1, 1] = np.NaN
plot_image(gauss, ax=ax1, scale='log')
# Run with all-NaN image:
gauss[:, :] = np.NaN
plot_image(gauss, ax=ax2, cbar='right')
return fig
def test_known_star(SHARED_INPUT_DIR, corrector, starid, cadence, var_goal, rms_goal, ptp_goal):
""" Check that the ensemble returns values that are reasonable and within expected bounds """
# All stars we check here come from the same sector and camera.
# Define these here for the future where we may test on other combinations of these:
sector = 1
camera = 1
__dir__ = os.path.abspath(os.path.dirname(__file__))
logger = logging.getLogger(__name__)
logger.info("-------------------------------------------------------------")
logger.info("CORRECTOR = %s, SECTOR=%d, CADENCE=%s, STARID=%d", corrector, sector, cadence, starid)
# All stars are from the same CCD, find the task for it:
with corrections.TaskManager(SHARED_INPUT_DIR) as tm:
task = tm.get_task(starid=starid, sector=sector, camera=camera, cadence=cadence)
# Check that task was actually found:
assert task is not None, "Task could not be found"
# Load lightcurve that will also be plotted together with the result:
# This lightcurve is of the same objects, at a state where it was deemed that the
# corrections were doing a good job.
compare_lc_path = os.path.join(__dir__, 'compare', f'compare-{corrector}-s{sector:04d}-c{cadence:04d}-tic{starid:011d}.ecsv.gz')
compare_lc = None
if os.path.isfile(compare_lc_path):
compare_lc = Table.read(compare_lc_path, format='ascii.ecsv')
else:
warnings.warn("Comparison data does not exist: " + compare_lc_path)
# Initiate the class
CorrClass = corrections.corrclass(corrector)
with tempfile.TemporaryDirectory() as tmpdir:
with CorrClass(SHARED_INPUT_DIR, plot=True) as corr:
# Check basic parameters of object (from BaseCorrector):
assert corr.input_folder == SHARED_INPUT_DIR, "Incorrect input folder"
assert corr.plot, "Plot parameter passed appropriately"
assert os.path.isdir(corr.data_folder), "DATA_FOLDER doesn't exist"
# Load the input lightcurve:
inlc = corr.load_lightcurve(task)
# Print input lightcurve properties:
print( inlc.show_properties() )
assert inlc.sector == sector
assert inlc.camera == camera
# Run correction:
tmplc = inlc.copy()
outlc, status = corr.do_correction(tmplc)
# Check status
assert outlc is not None, "Correction fails"
assert isinstance(outlc, TessLightCurve), "Should return TessLightCurve object"
assert isinstance(status, corrections.STATUS), "Should return a STATUS object"
assert status in (corrections.STATUS.OK, corrections.STATUS.WARNING), "STATUS was not set appropriately"
# Print output lightcurve properties:
print( outlc.show_properties() )
# Save the lightcurve to FITS file to be tested later on:
save_file = corr.save_lightcurve(outlc, output_folder=tmpdir)
# Check contents
assert len(outlc) == len(inlc), "Input flux ix different length to output flux"
assert isinstance(outlc.flux, np.ndarray), "FLUX is not a ndarray"
assert isinstance(outlc.flux_err, np.ndarray), "FLUX_ERR is not a ndarray"
assert isinstance(outlc.quality, np.ndarray), "QUALITY is not a ndarray"
assert outlc.flux.dtype.type is inlc.flux.dtype.type, "FLUX changes dtype"
assert outlc.flux_err.dtype.type is inlc.flux_err.dtype.type, "FLUX_ERR changes dtype"
assert outlc.quality.dtype.type is inlc.quality.dtype.type, "QUALITY changes dtype"
assert outlc.flux.shape == inlc.flux.shape, "FLUX changes shape"
assert outlc.flux_err.shape == inlc.flux_err.shape, "FLUX_ERR changes shape"
assert outlc.quality.shape == inlc.quality.shape, "QUALITY changes shape"
# Plot output lightcurves:
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, squeeze=True, figsize=[10, 10])
ax1.plot(inlc.time, inlc.flux, lw=0.5)
ax1.set_title(f"{corrector} - Sector {sector:d} - {cadence}s - TIC {starid:d}")
if compare_lc:
ax2.plot(compare_lc['time'], compare_lc['flux'], label='Compare', lw=0.5)
ax3.axhline(0, lw=0.5, ls=':', color='0.7')
ax3.plot(outlc.time, outlc.flux - compare_lc['flux'], lw=0.5)
ax2.plot(outlc.time, outlc.flux, label='New', lw=0.5)
ax1.set_ylabel('Flux [e/s]')
ax1.minorticks_on()
ax2.set_ylabel('Relative Flux [ppm]')
ax2.minorticks_on()
ax2.legend()
ax3.set_ylabel('New - Compare [ppm]')
ax3.set_xlabel('Time [TBJD]')
ax3.minorticks_on()
fig.savefig(os.path.join(__dir__, f'test-{corrector}-s{sector:04d}-c{cadence:04d}-tic{starid:011d}.png'), bbox_inches='tight')
plt.close(fig)
# Check things that are allowed to change:
assert all(outlc.flux != inlc.flux), "Input and output flux are identical."
assert not np.any(np.isinf(outlc.flux)), "FLUX contains Infinite"
assert not np.any(np.isinf(outlc.flux_err)), "FLUX_ERR contains Infinite"
assert np.sum(np.isnan(outlc.flux)) < 0.5*len(outlc), "More than half the lightcurve is NaN"
assert allnan(outlc.flux_err[np.isnan(outlc.flux)]), "FLUX_ERR should be NaN where FLUX is"
# TODO: Check that quality hasn't changed in ways that are not allowed:
# - Only values defined in CorrectorQualityFlags
# - No removal of flags already set
assert all(outlc.quality >= 0)
assert all(outlc.quality <= 128)
assert all(outlc.quality >= inlc.quality)
# Things that shouldn't chance from the corrections:
assert outlc.targetid == inlc.targetid, "TARGETID has changed"
assert outlc.label == inlc.label, "LABEL has changed"
assert outlc.sector == inlc.sector, "SECTOR has changed"
assert outlc.camera == inlc.camera, "CAMERA has changed"
assert outlc.ccd == inlc.ccd, "CCD has changed"
assert outlc.quality_bitmask == inlc.quality_bitmask, "QUALITY_BITMASK has changed"
assert outlc.ra == inlc.ra, "RA has changed"
assert outlc.dec == inlc.dec, "DEC has changed"
assert outlc.mission == 'TESS', "MISSION has changed"
assert outlc.time_format == 'btjd', "TIME_FORMAT has changed"
assert outlc.time_scale == 'tdb', "TIME_SCALE has changed"
assert_array_equal(outlc.time, inlc.time, "TIME has changed")
assert_array_equal(outlc.timecorr, inlc.timecorr, "TIMECORR has changed")
assert_array_equal(outlc.cadenceno, inlc.cadenceno, "CADENCENO has changed")
assert_array_equal(outlc.pixel_quality, inlc.pixel_quality, "PIXEL_QUALITY has changed")
assert_array_equal(outlc.centroid_col, inlc.centroid_col, "CENTROID_COL has changed")
assert_array_equal(outlc.centroid_row, inlc.centroid_row, "CENTROID_ROW has changed")
# Check metadata
assert tmplc.meta == inlc.meta, "Correction changed METADATA in-place"
assert outlc.meta['task'] == inlc.meta['task'], "Metadata is incomplete"
assert isinstance(outlc.meta['additional_headers'], fits.Header)
# Check performance metrics:
#logger.warning("VAR: %e", nanvar(outlc.flux))
if var_goal is not None:
var_in = nanvar(inlc.flux)
var_out = nanvar(outlc.flux)
var_diff = np.abs(var_out - var_goal) / var_goal
logger.info("VAR: %f - %f - %f", var_in, var_out, var_diff)
assert_array_less(var_diff, 0.05, "VARIANCE changed outside interval")
#logger.warning("RMS: %e", rms_timescale(outlc))
if rms_goal is not None:
rms_in = rms_timescale(inlc)
rms_out = rms_timescale(outlc)
rms_diff = np.abs(rms_out - rms_goal) / rms_goal
logger.info("RMS: %f - %f - %f", rms_in, rms_out, rms_diff)
assert_array_less(rms_diff, 0.05, "RMS changed outside interval")
#logger.warning("PTP: %e", ptp(outlc))
if ptp_goal is not None:
ptp_in = ptp(inlc)
ptp_out = ptp(outlc)
ptp_diff = np.abs(ptp_out - ptp_goal) / ptp_goal
logger.info("PTP: %f - %f - %f", ptp_in, ptp_out, ptp_diff)
assert_array_less(ptp_diff, 0.05, "PTP changed outside interval")
# Check FITS file:
with fits.open(os.path.join(tmpdir, save_file), mode='readonly') as hdu:
# Lightcurve FITS table:
fitslc = hdu['LIGHTCURVE'].data
hdr = hdu['LIGHTCURVE'].header
# Simple checks of header values:
assert hdu[0].header['TICID'] == starid
# Checks of things in FITS table that should not have changed at all:
assert_array_equal(fitslc['TIME'], inlc.time, "FITS: TIME has changed")
assert_array_equal(fitslc['TIMECORR'], inlc.timecorr, "FITS: TIMECORR has changed")
assert_array_equal(fitslc['CADENCENO'], inlc.cadenceno, "FITS: CADENCENO has changed")
assert_array_equal(fitslc['FLUX_RAW'], inlc.flux, "FITS: FLUX_RAW has changed")
assert_array_equal(fitslc['FLUX_RAW_ERR'], inlc.flux_err, "FITS: FLUX_RAW_ERR has changed")
assert_array_equal(fitslc['MOM_CENTR1'], inlc.centroid_col, "FITS: CENTROID_COL has changed")
assert_array_equal(fitslc['MOM_CENTR2'], inlc.centroid_row, "FITS: CENTROID_ROW has changed")
# Some things are allowed to change, but still within some requirements:
assert all(fitslc['FLUX_CORR'] != inlc.flux), "FITS: Input and output flux are identical."
assert np.sum(np.isnan(fitslc['FLUX_CORR'])) < 0.5*len(fitslc['TIME']), "FITS: More than half the lightcurve is NaN"
assert allnan(fitslc['FLUX_CORR_ERR'][np.isnan(fitslc['FLUX_CORR'])]), "FITS: FLUX_ERR should be NaN where FLUX is"
if corrector == 'ensemble':
# Check special headers:
assert np.isfinite(hdr['ENS_MED']) and hdr['ENS_MED'] > 0
assert isinstance(hdr['ENS_NUM'], int) and hdr['ENS_NUM'] > 0
assert hdr['ENS_DLIM'] == 1.0
assert hdr['ENS_DREL'] == 10.0
assert hdr['ENS_RLIM'] == 0.4
# Special extension for ensemble:
tic = hdu['ENSEMBLE'].data['TIC']
bzeta = hdu['ENSEMBLE'].data['BZETA']
assert len(tic) == len(bzeta)
assert len(np.unique(tic)) == len(tic), "TIC numbers in ENSEMBLE table are not unique"
assert len(tic) == hdr['ENS_NUM'], "Not the same number of targets in ENSEMBLE table as specified in header"
elif corrector == 'cbv':
# Check special headers:
assert isinstance(hdr['CBV_NUM'], int) and hdr['CBV_NUM'] > 0
# Check coefficients:
for k in range(0, hdr['CBV_NUM']+1):
assert np.isfinite(hdr['CBV_C%d' % k])
for k in range(1, hdr['CBV_NUM']+1):
assert np.isfinite(hdr['CBVS_C%d' % k])
# Check that no other coefficients are present
assert 'CBV_C%d' % (hdr['CBV_NUM']+1) not in hdr
assert 'CBVS_C%d' % (hdr['CBV_NUM']+1) not in hdr
elif corrector == 'kasoc_filter':
# Check special headers:
assert hdr['KF_POSS'] == 'None'
assert np.isfinite(hdr['KF_LONG']) and hdr['KF_LONG'] > 0
assert np.isfinite(hdr['KF_SHORT']) and hdr['KF_SHORT'] > 0
assert hdr['KF_SCLIP'] == 4.5
assert hdr['KF_TCLIP'] == 5.0
assert hdr['KF_TWDTH'] == 1.0
assert hdr['KF_PSMTH'] == 200
assert isinstance(hdr['NUM_PER'], int) and hdr['NUM_PER'] >= 0
for k in range(1, hdr['NUM_PER']+1):
assert np.isfinite(hdr['PER_%d' % k]) and hdr['PER_%d' % k] > 0
# Check that no other periods are present
assert 'PER_%d' % (hdr['NUM_PER'] + 1) not in hdr
# Test that the Gzip FITS file has the correct uncompressed file name, by simply
# decompressing the Gzip file, asking to keep the original file name.
# This uses the system GZIP utility, since there doesn't seem to be a way to do this
# through the Python gzip module:
fpath = os.path.join(tmpdir, save_file)
fpath_uncompressed = fpath.replace('.fits.gz', '.fits')
assert not os.path.exists(fpath_uncompressed), "Uncompressed file already exists"
gzip_output = subprocess.check_output(['gzip', '-dkNv', os.path.basename(fpath)],
cwd=os.path.dirname(fpath),
stderr=subprocess.STDOUT,
encoding='utf8')
print("Gzip output:")
print(gzip_output)
assert os.path.isfile(fpath_uncompressed), "Incorrect uncompressed file name"
# Just see if we can in fact also open the uncompressed FITS file and get a simple header:
with fits.open(fpath_uncompressed, mode='readonly') as hdu:
assert hdu[0].header['TICID'] == starid