import gc import os import argparse import functools import numpy as np import pandas as pd from astropy.stats import sigma_clipped_stats import matplotlib.pyplot as plt QUARTER_THRESHOLD = 1.5 def with_gc_collect(func): @functools.wraps(func) def wrapper_decorator(*args, **kwargs): value = func(*args, **kwargs) gc.collect() return value return wrapper_decorator def get_segments(time: np.ndarray, gap_size: float) -> np.ndarray: args, = np.where(np.diff(time) > gap_size) args = args + 1 segments = np.zeros((len(args) + 1, 2), dtype='int') segments[1:, 0] = args segments[:-1, 1] = args segments[-1, 1] = len(time) return segments def bin_lightcurve(time: np.ndarray, flux: np.ndarray, flux_err: np.ndarray, bin_size: float = 12, method: str = 'points' ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """ Re-bin lightcurve. """ if method == 'points': bin_idx = np.arange(len(time)) // bin_size elif method == 'window': nbins = np.ceil(np.ptp(time)/bin_size).astype('int') bin_edges = np.amin(time) + bin_size*np.arange(nbins + 1) bin_idx = np.searchsorted(bin_edges, time, side='right') else: raise ValueError(f"Unknown binning method: {method}.") weights = 1/flux_err**2 points = np.bincount(bin_idx) time_sum = np.bincount(bin_idx, weights=time) weights_sum = np.bincount(bin_idx, weights=weights) weights_flux_sum = np.bincount(bin_idx, weights=weights*flux) # Remove empty bins. mask = points > 0 points = points[mask] time_sum = time_sum[mask] weights_sum = weights_sum[mask] weights_flux_sum = weights_flux_sum[mask] time = time_sum/points flux = weights_flux_sum/weights_sum flux_err = np.sqrt(1/weights_sum) # plt.plot(lc_data['time'][::10], lc_data['flux'][::10], '.') # plt.errorbar(time, flux, yerr=flux_err, ls='none') # plt.show() # plt.close() return time, flux, flux_err, points def get_binned_lightcurve(lc: pd.DataFrame, bin_width_mins: float) -> pd.DataFrame: # Compute the binned lightcurve. result = bin_lightcurve(lc['time'], lc['flux'], lc['flux_err'], bin_size=bin_width_mins * 60, method='window') time_bin, flux_bin, flux_err_bin, points = result # Create a new data frame. lc_bin = pd.DataFrame(time_bin, columns=['time']) lc_bin['flux'] = flux_bin lc_bin['flux_err'] = flux_err_bin lc_bin['points'] = points return lc_bin def get_lightcurve(lc_master: pd.DataFrame, transit: str, components: list[str]) -> pd.DataFrame: random = 0 if 'random' in components: random = lc_master['random'] systematics = 0 if 'systematics' in components: systematics = lc_master['systematics'] granulation = 0 if 'granulation' in components: granulation = lc_master['granulation'] oscillations = 0 if 'oscillations' in components: oscillations = lc_master['oscillations'] spots = 1 if 'spots' in components: spots = lc_master['spots'] if transit is None: transit_signal = 1 else: transit_signal = lc_master[f'transit_{transit}'] # Compute the total flux from noise and signal. noise = (1 + (granulation + oscillations) * 1e-6) * (1 + (random + systematics) * 1e-6) flux = noise * spots * transit_signal # Determine the lightcurve quarters. quarters = get_segments(lc_master['time'] / (24 * 3600), QUARTER_THRESHOLD) # Compute the flux uncertainty on each quarter. flux_err = np.zeros_like(flux) for imin, imax in quarters: _, _, stdev = sigma_clipped_stats(np.diff(flux[imin:imax])) flux_err[imin:imax] = stdev / np.sqrt(2) # Normalize the lightcurve. norm = np.median(flux) flux = flux / norm flux_err = flux_err / norm lc = pd.DataFrame(lc_master['time'], columns=['time']) lc['flux'] = flux lc['flux_err'] = flux_err lc['flag'] = lc_master['flag'] return lc def get_lightcurves(filelist: list[str], output_dir: str, transit: str, components: list[str], bin_width_mins: float ) -> None: if transit not in [None, 'inner', 'outer', 'inner_ecc', 'outer_ecc']: raise ValueError(f"Unknown transit signal type: {transit}, must be 'inner', 'outer', 'inner_ecc' or 'outer_ecc'.") for item in components: if item not in [None, 'random', 'systematics', 'granulation', 'oscillations', 'spots']: raise ValueError(f"Unknown signal type: {item}, only 'granulation', 'oscillations' or 'spots' are valid.") # Create the output directory. os.makedirs(output_dir, exist_ok=True) for filename in filelist: # Check if the output file exists. basename = os.path.basename(filename) outfile = os.path.join(output_dir, basename) if os.path.exists(outfile): print(f'Output file {outfile} exists, skipping...') continue # Read the lightcurve master file. lc_master = pd.read_feather(filename) # Generate a lightcurve containing the requested signals. lc = get_lightcurve(lc_master, transit, components) # If requested bin the lighcurve. if bin_width_mins is not None: lc = get_binned_lightcurve(lc, bin_width_mins) # Write the generated lightcurve to file. lc.to_feather(outfile) return def plot_segments(x, y, segments, *args, **kwargs): for imin, imax in segments: plt.plot(x[imin:imax], y[imin:imax], *args, **kwargs) return @with_gc_collect def plot_master_file(filename): head, tail = os.path.split(filename) figdir = os.path.join(head, 'figures') os.makedirs(figdir, exist_ok=True) figfile = os.path.join(figdir, tail.replace('.ftr', '.png')) df = pd.read_feather(filename) time = df['time'] / (24 * 3600) segments = get_segments(time, QUARTER_THRESHOLD) plt.figure(figsize=(13, 21)) ax = plt.subplot(811) plot_segments(time, df['granulation'], segments) plt.ylabel('granulation [ppm]') plt.subplot(812, sharex=ax) plot_segments(time, df['oscillations'], segments) plt.ylabel('oscillations [ppm]') plt.subplot(813, sharex=ax) plot_segments(time, df['systematics'], segments) plt.ylabel('systematics [ppm]') plt.subplot(814, sharex=ax) plot_segments(time, df['random'], segments, '.') plt.ylabel('random [ppm]') plt.subplot(815, sharex=ax) plot_segments(time, df['noise'], segments, '.') plt.ylabel('combined noise [ppm]') plt.subplot(816, sharex=ax) plot_segments(time, (df['spots'] - 1)*1e6, segments) plt.ylabel('spots') plt.subplot(817, sharex=ax) plot_segments(time, (df['transit_inner'] - 1) * 1e6, segments) plt.ylabel('transit innner [ppm]') plt.subplot(818, sharex=ax) plot_segments(time, (df['transit_outer'] - 1) * 1e6, segments) plt.ylabel('transit outer [ppm]') plt.xlabel('Time [days]') plt.xlim(0, np.amax(time)) plt.tight_layout() plt.savefig(figfile, dpi=300) # plt.show() plt.close() return def generate_func(args): get_lightcurves(args.filelist, args.output_dir, args.transit, args.components, args.bin_width_mins) def plot_func(args): for filename in args.filelist: plot_master_file(filename) def main(): parser = argparse.ArgumentParser(prog='plato-lightcurve', description='process the master lightcurve files', formatter_class=argparse.ArgumentDefaultsHelpFormatter) subparsers = parser.add_subparsers(title='subcommands', required=True, description='valid subcommands', help='list of valid subcommands', dest='subcommand') gen = subparsers.add_parser('generate', help='generate lightcurve files', formatter_class=argparse.ArgumentDefaultsHelpFormatter) gen.add_argument('filelist', type=str, nargs='+', help='the master files to process') gen.add_argument('output_dir', type=str, help='the location to write the output files to, can not be an existing directory') gen.add_argument('--components', type=str, nargs='+', default=None, choices=['random', 'systematics', 'granulation', 'oscillations', 'spots'], help='the non-transit signal(s) to include') gen.add_argument('--transit', type=str, default=None, choices=['inner', 'outer', 'inner_ecc', 'outer_ecc'], help='the transit signal, if any, to include') gen.add_argument('--bin-width-mins', type=float, default=None, help='the width of the time bins, if any, in minutes') gen.set_defaults(func=generate_func) plot = subparsers.add_parser('plot', help='plot master lightcurves', formatter_class=argparse.ArgumentDefaultsHelpFormatter) plot.add_argument('filelist', type=str, nargs='+', help='the master files to process') plot.set_defaults(func=plot_func) # Parse the arguments and call the function. args = parser.parse_args() args.func(args) return if __name__ == '__main__': main()