<filename>lightkurve_ext.py<gh_stars>1-10
"""
Convenience helpers for `lightkurve` package.
"""

import os
import logging
import math
import json
import re
import warnings
from collections import OrderedDict

from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.table import Table
from astropy.time import Time
import astropy.units as u

import numpy as np
from scipy.interpolate import UnivariateSpline

from IPython.display import display, HTML

import lightkurve as lk

import asyncio_compat

log = logging.getLogger(__name__)


def of_sector(lcf_coll, sectorNum):
    res_list = of_sectors(lcf_coll, sectorNum)
    return res_list[0] if len(res_list) > 0 else None


def of_sectors(*args):
    lk_coll_or_sr = args[0]
    if len(args) == 1:
        # when no sectors are specified, return entire collection
        # For convenience: when a notebooks is modified such that
        # a user sometimes use a subset of sectors , and sometimes everything
        # the user can can still use of_sectors() wrapper regardless
        return lk_coll_or_sr
    sector_nums = args[1:]

    if hasattr(lk_coll_or_sr, "sector"):
        return lk_coll_or_sr[np.in1d(lk_coll_or_sr.sector, sector_nums)]
    elif hasattr(lk_coll_or_sr, "table") and lk_coll_or_sr.table["sequence_number"] is not None:
        return lk.SearchResult(lk_coll_or_sr.table[np.in1d(lk_coll_or_sr.table["sequence_number"], sector_nums)])
    else:
        raise TypeError(f"Unsupported type of collection: {type(lk_coll_or_sr)}")


def of_sector_n_around(lk_coll_or_sr, sector_num, num_additions=8):
    def get_sector_for_lk_coll(lk_coll):
        return lk_coll.sector

    def get_sector_for_sr(sr):
        return sr.table["sequence_number"]

    sector_accessor_func = None
    if hasattr(lk_coll_or_sr, "sector"):
        sector_accessor_func = get_sector_for_lk_coll
    elif hasattr(lk_coll_or_sr, "table") and lk_coll_or_sr.table["sequence_number"] is not None:
        sector_accessor_func = get_sector_for_sr
    else:
        raise TypeError(f"Unsupported type of collection: {type(lk_coll_or_sr)}")

    subset_slice = _get_slice_for_of_sector_n_around(
        lk_coll_or_sr,
        sector_accessor_func,
        sector_num,
        num_additions=num_additions,
    )
    if subset_slice is not None:
        return lk_coll_or_sr[subset_slice]
    else:
        return type(lk_coll_or_sr)([])


def _get_slice_for_of_sector_n_around(coll, sector_accessor_func, sector_num, num_additions):
    if sector_num not in sector_accessor_func(coll):
        return None

    idx = np.where(sector_accessor_func(coll) == sector_num)[0][0]
    # if num_additions is odd number, we add one to older sector
    start = max(idx - math.ceil(num_additions / 2), 0)
    end = min(idx + math.floor(num_additions / 2) + 1, len(coll))

    # case the start:end slice does not fill up the requested num_additions,
    # try to fill it up
    cur_slice_size = end - start - 1
    if cur_slice_size < num_additions:
        num_more_needed = num_additions - cur_slice_size
        if start > 0:
            start = max(start - num_more_needed, 0)
        else:
            end = min(end + num_more_needed, len(coll))

    return slice(start, end)


def of_2min_cadences(lcf_coll):
    """Return LightCurveFiles of short, typically 2-minute cadence, only.
    Primary use case is to filter out 20-second files.
    """
    filtered = [lcf for lcf in lcf_coll if "short" == estimate_cadence_type(lcf)]
    return lk.LightCurveCollection(filtered)


def estimate_cadence(lc):
    """Estimate the cadence of a lightcurve by returning the median of a sample"""
    return np.nanmedian(np.diff(lc.time[:100].value))


def map_cadence_type(cadence_in_days):
    long_minimum = 9.9 / 60 / 24  # 10 minutes in days, with some margin of error
    short_minimum = 0.9 / 60 / 24  # 1 minute in days, with some margin of error
    if cadence_in_days is None:
        return None
    if cadence_in_days >= long_minimum:
        return "long"
    if cadence_in_days >= short_minimum:
        return "short"
    return "fast"


def estimate_cadence_type(lc):
    """Estimate the type of cadence to be one of long, short, or fast.
    The definition is the same as ``exptime`` in `lightkurve.search_lightcurve()`.
    """
    return map_cadence_type(estimate_cadence(lc))


def of_tic(lcf_coll, tic):
    """Return LightCurveFiles of the given TIC.

    Useful in case the default MAST result returned nearby targets.
    """
    filtered = [lcf for lcf in lcf_coll if lcf.meta.get("TICID", None) == tic]
    return lk.LightCurveCollection(filtered)


def estimate_object_radius_in_r_jupiter(lc, depth):
    """Return a back of envelope estimate of a companion object's radius."""
    R_JUPITER_IN_R_SUN = 71492 / 695700

    r_star = lc.meta.get("RADIUS")  # assumed to be in R_sun
    if r_star is None or r_star < 0 or depth <= 0:
        return None  # cannot estimate
    r_obj = math.sqrt(r_star * r_star * depth)
    r_obj_in_r_jupiter = r_obj / R_JUPITER_IN_R_SUN
    return r_obj_in_r_jupiter


def download_lightcurves_of_tic_with_priority(tic, download_filter_func=None, download_dir=None):
    """For a given TIC, download lightcurves across all sectors.
    For each sector, download one based on pre-set priority.
    """

    sr_unfiltered = lk.search_lightcurve(f"TIC{tic}", mission="TESS")
    if len(sr_unfiltered) < 1:
        print(f"WARNING: no result found for TIC {tic}")
        return None, None, None

    sr_unfiltered = sr_unfiltered[sr_unfiltered.target_name == str(tic)]  # in case we get some other nearby TICs

    # filter out HLSPs not supported by lightkurve yet
    sr = sr_unfiltered[sr_unfiltered.author != "DIAMANTE"]
    if len(sr) < len(sr_unfiltered):
        print("Note: there are products not supported by Lightkurve, which are excluded from download.")

    # for each sector, filter based on the given priority.
    # - note: prefer QLP over TESS-SPOC because QLP is detrended, with multiple apertures within 1 file
    sr = filter_by_priority(
        sr,
        author_priority=["SPOC", "QLP", "TESS-SPOC"],
        exptime_priority=["short", "long", "fast"],
    )
    num_filtered = len(sr_unfiltered) - len(sr)
    num_fast = len(sr_unfiltered[sr_unfiltered.exptime < 60 * u.second])
    if num_filtered > 0:
        msg = f"{num_filtered} rows filtered"
        if num_fast > 0:
            msg = msg + f" ; {num_fast} fast (20secs) products."
        print(msg)

    display(sr)

    # let caller to optionally further restrict a subset to be downloaded
    sr_to_download = sr
    if download_filter_func is not None:
        sr_to_download = download_filter_func(sr)
        if len(sr_to_download) < len(sr):
            display(
                HTML(
                    """<font style="background-color: yellow;">Note</font>:
SearchResult is further filtered - only a subset will be downloaded."""
                )
            )

    lcf_coll = sr_to_download.download_all(download_dir=download_dir)

    if lcf_coll is not None and len(lcf_coll) > 0:
        print(f"TIC {tic} \t#sectors: {len(lcf_coll)} ; {lcf_coll[0].meta['SECTOR']} - {lcf_coll[-1].meta['SECTOR']}")
        print(
            (
                f"   sector {lcf_coll[-1].meta['SECTOR']}: \t"
                f"camera = {lcf_coll[-1].meta['CAMERA']} ; ccd = {lcf_coll[-1].meta['CCD']}"
            )
        )
    else:
        print(f"TIC {tic}: no data")

    return lcf_coll, sr, sr_unfiltered


def download_lightcurve(
    target,
    mission=("Kepler", "K2", "TESS"),
    exptime="short",
    author="SPOC",
    download_dir=None,
    use_cache="yes",
    display_search_result=True,
):
    """
    Wraps `lightkurve.search_lightcurve()` and the
    subsequent `lightkurve.search.SearchResult.download_all()` calls,
    with the option of caching, so that for a given search,
    if the the result has been downloaded, the cache will be used.

    The parameters all propagate to the underlying `search_lightcurvefile()`
    and `download_all()` calls. The lone exception is `use_cache`.

    Parameters
    ----------
    use_cache : str, must be one of 'yes', or 'no'\n
        OPEN: an option of 'fallback': cache will be used when offline.\n
        OPEN: for now, actual download lightcurve cache will still be used if
        available irrespective of the settings.

    Returns
    -------
    collection : `~lightkurve.collections.Collection` object
        Returns a `~lightkurve.collections.LightCurveCollection`
        containing all lightcurve files that match the criteria
    """

    if use_cache == "no":
        return _search_and_cache(target, mission, exptime, author, download_dir, display_search_result)
    if use_cache == "yes":
        result_file_ids = _load_from_cache_if_any(target, mission, download_dir)
        if result_file_ids is not None:
            result_files = list(map(lambda e: f"{download_dir}/mastDownload/{e}", result_file_ids))
            return lk.collections.LightCurveCollection(list(map(lambda f: lk.read(f), result_files)))
        # else
        return _search_and_cache(target, mission, exptime, author, download_dir, display_search_result)
    # else
    raise ValueError("invalid value for argument use_cache")


# Private helpers for `download_lightcurvefiles`


def _search_and_cache(target, mission, exptime, author, download_dir, display_search_result):
    search_res = lk.search_lightcurve(target=target, mission=mission, exptime=exptime, author=author)
    if len(search_res) < 1:
        return None
    if display_search_result:
        _display_search_result(search_res)
    _cache_search_result_product_identifiers(search_res, download_dir, target, mission)
    return search_res.download_all(quality_bitmask="default", download_dir=download_dir)


def _display_search_result(search_res):
    from IPython.core.display import display

    tab = search_res.table
    # move useful columns to the front
    preferred_cols = ["proposal_id", "target_name", "sequence_number", "t_exptime"]
    colnames_reordered = preferred_cols + [c for c in tab.colnames if c not in preferred_cols]
    display(tab[colnames_reordered])


def _load_from_cache_if_any(target, mission, download_dir):
    key = _get_cache_key(target, mission)
    return _load_search_result_product_identifiers(download_dir, key)


def _cache_search_result_product_identifiers(search_res, download_dir, target, mission):
    key = _get_cache_key(target, mission)
    identifiers = _to_product_identifiers(search_res)
    _save_search_result_product_identifiers(identifiers, download_dir, key)
    return key


def _get_search_result_cache_dir(download_dir):
    # TODO: handle download_dir is None (defaults)
    cache_dir = f"{download_dir}/mastQueries"

    if os.path.isdir(cache_dir):
        return cache_dir

    # else it doesn't exist, make a new cache directory
    try:
        os.mkdir(cache_dir)
    # downloads locally if OS error occurs
    except OSError:
        log.warning(
            "Warning: unable to create {}. "
            "Cache MAST query results to the current "
            "working directory instead.".format(cache_dir)
        )
        cache_dir = "."
    return cache_dir


def _get_cache_key(target, mission):
    # TODO: handle cases the generated key is not a valid filename
    return f"{target}_{mission}_ids"


def _to_product_identifiers(search_res):
    """
    Returns
    -------
    A list of str, constructed from `(obs_collection, obs_id, productFilename)` tuples, that can
    identify cached lightcurve file,s if any.
    """
    return list(
        map(
            lambda e: e["obs_collection"] + "/" + e["obs_id"] + "/" + e["productFilename"],
            search_res.table,
        )
    )


def _save_search_result_product_identifiers(identifiers, download_dir, key):
    resolved_cache_dir = _get_search_result_cache_dir(download_dir)
    filepath = f"{resolved_cache_dir}/{key}.json"
    fp = open(filepath, "w+")
    json.dump(identifiers, fp)
    return filepath


def _load_search_result_product_identifiers(download_dir, key):
    resolved_cache_dir = _get_search_result_cache_dir(download_dir)
    filepath = f"{resolved_cache_dir}/{key}.json"
    try:
        fp = open(filepath, "r")
        return json.load(fp)
    except OSError as err:
        # errno == 2: file not found, typical case of cache miss
        # errno != 2: unexpected error, log a warning
        if err.errno != 2:
            log.warning("Unexpected OSError in retrieving cached search result: {}".format(err))
        return None


def filter_by_priority(
    sr,
    author_priority=["SPOC", "TESS-SPOC", "QLP"],
    exptime_priority=["short", "long", "fast"],
):
    author_sort_keys = {}
    for idx, author in enumerate(author_priority):
        author_sort_keys[author] = idx + 1

    exptime_sort_keys = {}
    for idx, exptime in enumerate(exptime_priority):
        exptime_sort_keys[exptime] = idx + 1

    def calc_filter_priority(row):
        # Overall priority key is in the form of <author_key><exptime_key>, e.g., 101
        # - "01" is the exptime_key
        # - the leading "1" is the author_key, given it is the primary one
        author_default = max(dict(author_sort_keys).values()) + 1
        author_key = author_sort_keys.get(row["author"], author_default) * 100

        # secondary priority
        exptime_default = max(dict(exptime_sort_keys).values()) + 1
        exptime_key = exptime_sort_keys.get(map_cadence_type(row["exptime"] / 60 / 60 / 24), exptime_default)
        return author_key + exptime_key

    sr.table["_filter_priority"] = [calc_filter_priority(r) for r in sr.table]

    # A temporary table that sorts the table by the priority
    sorted_t = sr.table.copy()
    sorted_t.sort(["mission", "_filter_priority"])

    # create an empty table for results, with the same set of columns
    res_t = sr.table[np.zeros(len(sr), dtype=bool)].copy()

    # for each mission (e.g., TESS Sector 01), select a row based on specified priority
    # - select the first row given the table has been sorted by priority
    uniq_missions = list(OrderedDict.fromkeys(sorted_t["mission"]))
    for m in uniq_missions:
        mission_t = sorted_t[sorted_t["mission"] == m]
        # OPEN: if for a given mission, the only row available is not listed in the priorities,
        # the logic still add a row to the result.
        # We might want it to be an option specified by the user.
        res_t.add_row(mission_t[0])

    return lk.SearchResult(table=res_t)


# Download TPF asynchronously


def search_and_download_tpf(*args, **kwargs):
    """Search and Download a TPFs.

    All parameters are passed on ``search_targetpixelfile()``,
    with the exception of ``download_dir`` and ``quality_bitmask``,
    which are passed to ``download_all()`
    """

    # extract download_all() parameters
    download_dir = kwargs.pop("download_dir", None)
    quality_bitmask = kwargs.pop("quality_bitmask", None)
    sr = lk.search_targetpixelfile(*args, **kwargs)  # pass the rest of the argument to search_targetpixelfile
    tpf_coll = sr.download_all(download_dir=download_dir, quality_bitmask=quality_bitmask)
    return tpf_coll, sr


def create_download_tpf_task(*args, **kwargs):
    return asyncio_compat.create_background_task(search_and_download_tpf, *args, **kwargs)


#
# Other misc. extensions
#
def get_bkg_lightcurve(lcf):
    """Returns the background flux, i.e., ``SAP_BKG`` in the file"""
    lc = lcf.copy()
    lc["flux"] = lc["sap_bkg"]
    lc["flux_err"] = lc["sap_bkg_err"]
    lc.label = lc.label + " BKG"
    return lc


def _do_create_quality_issues_mask(quality, flux, flags_included=0b0101001010111111):
    """Returns a boolean array which flags cadences with *issues*.

    The default `flags_included` is a TESS default, based on
    https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags
    """
    if np.issubdtype(quality.dtype, np.integer):
        return np.logical_and(quality & flags_included, np.isfinite(flux))
    else:
        # quality column is not an integer, probably a non-standard product
        return np.zeros_like(quality, dtype=bool)


def create_quality_issues_mask(lc, flags_included=0b0101001010111111):
    """Returns a boolean array which flags cadences with *issues*.

    The default `flags_included` is a TESS default, based on
    https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags
    """

    # use sap_flux when available (it may not be there in some HLSP)
    # we prefer sap_flux over pdcsap_flux as
    # pdcsap_flux is more likely to be NaN (due to exclusion by quality flags)
    if "sap_flux" in lc.colnames:
        flux = lc["sap_flux"]
    else:
        flux = lc.flux

    return _do_create_quality_issues_mask(lc.quality, flux)


def _get_n_truncate_fits_data(lc, before, after, return_columns, return_mask=False):
    with fits.open(lc.filename) as hdu:
        time = hdu[1].data["TIME"]
        mask = (time >= before) & (time < after)
        res = dict()
        for col in return_columns:
            res[col] = hdu[1].data[col][mask]
        if return_mask:
            return res, mask
        else:
            return res


def list_times_w_quality_issues(lc, include_excluded_cadences=False):
    if not include_excluded_cadences:
        mask = create_quality_issues_mask(lc)
        return lc.time[mask], lc.quality[mask]
    else:
        # case we want cadences that have been excluded in the lc object
        # use the underlying fits file

        flux_colname = lc.meta.get("FLUX_ORIGIN", "sap_flux")
        if flux_colname == "pdcsap_flux":
            flux_colname = "sap_flux"  # pdcsap_flux would not have values in excluded cadences, defeating the purpose

        # data is truncated if the given lc is truncated
        # TODO: cadences excluded before lc.time.min() or after lc.time.max() will still be missing.
        data = _get_n_truncate_fits_data(lc, lc.time.min().value, lc.time.max().value, ["time", "quality", flux_colname])
        mask = _do_create_quality_issues_mask(data["quality"], data[flux_colname])
        return data["time"][mask], data["quality"][mask]


def list_transit_times(t0, period, steps_or_num_transits=range(0, 10), return_string=False):
    """List the transit times based on the supplied transit parameters"""
    if isinstance(steps_or_num_transits, int):
        steps = range(0, steps_or_num_transits)
    else:
        steps = steps_or_num_transits
    times = [t0 + period * i for i in steps]
    if return_string:
        return ",".join(map(str, times))
    else:
        return times


def get_segment_times_idx(times, break_tolerance=5):
    """Segment the input array of times into segments due to data gaps. Return the indices of the segments.

    The minimal gap size is determined by `break_tolerance`.

    The logic is adapted from `LightCurve.flatten`
    """
    if hasattr(times, "value"):  # convert astropy Time to raw values if needed
        times = times.value
    dt = times[1:] - times[0:-1]
    with warnings.catch_warnings():  # Ignore warnings due to NaNs
        warnings.simplefilter("ignore", RuntimeWarning)
        cut = np.where(dt > break_tolerance * np.nanmedian(dt))[0] + 1
    low = np.append([0], cut)
    high = np.append(cut, len(times))
    return (low, high)


def get_segment_times(times, **kwargs):
    if hasattr(times, "value"):  # convert astropy Time to raw values if needed
        times = times.value
    low, high = get_segment_times_idx(times, **kwargs)
    # add a small 1e-10 to end so that the end time is exclusive (follow convention in range)
    return [(times[lo], times[hi - 1] + 1e-10) for lo, hi in zip(low, high)]


def get_transit_times_in_range(t0, period, start, end):
    t_start = t0 + math.ceil((start - t0) / period) * period
    num_t = math.ceil((end - t_start) / period)
    return [t_start + period * i for i in range(num_t)]


def get_transit_times_in_lc(lc, t0, period, return_string=False, **kwargs):
    """Get the transit times with observations of the given lightcurve, based on the supplied transit parameters.

    The method will exclude the times where there is no observation due to data gaps.
    """

    lc = lc.remove_nans()  # exclude cadences with no flux.
    # break up the times to exclude times in gap
    times_list = get_segment_times(lc.time, **kwargs)
    transit_times = []
    for start, end in times_list:
        transit_times.extend(get_transit_times_in_range(t0, period, start, end))
    if return_string:
        return ",".join(map(str, transit_times))
    else:
        return transit_times


def stitch(lcf_coll, **kwargs):
    """Wrapper over native stitch(), and tweak the metadata so that it behaves like a typical single-sector lightcurve."""

    def update_meta_if_exists_in(lc_src, keys):
        for key in keys:
            val = lc_src.meta.get(key, None)
            if val is not None:
                lc_stitched.meta[key] = val

    def safe_del_meta(key):
        if lc_stitched.meta.get(key, None) is not None:
            del lc_stitched.meta[key]

    lc_stitched = lcf_coll.stitch(**kwargs)

    # now update the metadata

    lc_stitched.meta["STITCHED"] = True

    # update observation start/stop dates
    update_meta_if_exists_in(lcf_coll[0], ("TSTART", "DATE-OBS"))
    update_meta_if_exists_in(lcf_coll[-1], ("TSTOP", "DATE-END"))

    # TODO: recalculate TELAPSE, LIVETIME, DEADC (which is LIVETIME / TELAPSE)

    safe_del_meta("FILENAME")  # don't associate it with a file anymore

    # record the sectors stitched and the associated metadata
    sector_list = [lc.meta.get("SECTOR") for lc in lcf_coll if lc.meta.get("SECTOR") is not None]
    meta_list = [lc.meta for lc in lcf_coll if lc.meta.get("SECTOR") is not None]
    if len(sector_list) > 0:
        lc_stitched.meta["SECTORS"] = sector_list
        meta_dict = dict()
        for sector, meta in zip(sector_list, meta_list):
            meta_dict[sector] = meta.copy()
        lc_stitched.meta["HEADERS_ORIGINAL"] = meta_dict

    return lc_stitched


def to_window_length_for_2min_cadence(length_day):
    """Helper for LightCurve.flatten().
    Return a `window_length` for the given number of days, assuming the data has 2-minute cadence."""
    res = math.floor(720 * length_day)
    if res % 2 == 0:
        res += 1  # savgol_filter window length must be odd number
    return res


# detrend using spline
# Based on:  https://github.com/barentsen/kepler-athenaeum-tutorial/blob/master/how-to-find-a-planet-tutorial.ipynb
def flatten_with_spline_normalized(lc, return_trend=False, **kwargs):
    lc = lc.remove_nans()
    spline = UnivariateSpline(lc.time, lc.flux, **kwargs)
    trend = spline(lc.time)
    # detrended = lc.flux - trend
    detrended_relative = 100 * ((lc.flux / trend) - 1) + 100  # in percentage
    lc_flattened = lc.copy()
    lc_flattened.flux = detrended_relative
    lc_flattened.flux_unit = "percent"
    if not return_trend:
        return lc_flattened
    else:
        lc_trend = lc.copy()
        lc_trend.flux = trend
        return (lc_flattened, lc_trend)


def _lksl_statistics(ts):
    """Compute LKSL Statistics of the given (time-series) values.
    Useful to compare the noises in a folded lightcurve.

    Based on https://arxiv.org/pdf/1901.00009.pdf , equation 4.
    See https://www.aanda.org/articles/aa/pdf/2002/17/aa2208.pdf for more information.
    """
    ts = ts[~np.isnan(ts)]

    vector_length_sq_sum = np.square(np.diff(ts)).sum()
    if len(ts) > 2:
        # to fully utilize the data by including the vector length between the last and first measurement
        # section 2 of Clarke, 2002
        vector_length_sq_sum += np.square(ts[-1] - ts[0])

    diff_from_mean_sq_sum = np.square(ts - np.mean(ts)).sum()
    if diff_from_mean_sq_sum == 0:  # to avoid boundary case that'd cause division by zero
        diff_from_mean_sq_sum = 1e-10

    return (vector_length_sq_sum / diff_from_mean_sq_sum) * (len(ts) - 1) / (2 * len(ts))


def lksl_statistics(lc, column="flux"):
    return _lksl_statistics(lc[column].value)


def tess_flux_to_mag(flux):
    """Convert flux from TESS observation to magnitude."""
    # Based on https://heasarc.gsfc.nasa.gov/docs/tess/observing-technical.html#saturation
    # TESS CCDs produce 15000 e/s for magnitude 10 light source

    if isinstance(flux, u.Quantity):
        flux_raw = flux.to(u.electron / u.second).value
    else:
        flux_raw = flux

    # np.log10 does not work on Quantity, unless it's dimensionless
    res_raw = 10 - 2.5 * np.log10((flux_raw / 15000))

    if isinstance(flux, u.Quantity):
        return res_raw * u.mag
    else:
        return res_raw


def mag_to_tess_flux(mag):
    """Convert magnitude to flux in TESS observation."""
    if isinstance(mag, u.Quantity):
        mag_raw = mag.to(u.mag).value
    else:
        mag_raw = mag

    flux_raw = (10 ** ((10 - mag_raw) / 2.5)) * 15000
    if isinstance(mag, u.Quantity):
        return flux_raw * ((u.electron / u.second))
    else:
        return flux_raw


def select_flux(lc, flux_cols):
    """Return a Lightcurve object with the named column as the flux column.

    flux_cols: either a column name (string), or a list of prioritized column names
    such that the first one that the lightcurve contains will be used.
    """

    def _to_lc_with_1flux(lc, flux_1col):
        flux_1col = flux_1col.lower()
        if "flux" == flux_1col:
            return lc
        elif flux_1col in lc.colnames:
            return lc.select_flux(flux_1col)
        else:
            return None

    if isinstance(flux_cols, str):
        flux_cols = [flux_cols]

    for flux_1col in flux_cols:
        res = _to_lc_with_1flux(lc, flux_1col)
        if res is not None:
            return res
    raise ValueError(f"'column {flux_cols}' not found")


#
# TargetPixelFile helpers
#


def truncate(tpf: lk.targetpixelfile.TargetPixelFile, before: float, after: float) -> lk.targetpixelfile.TargetPixelFile:
    return tpf[(tpf.time.value >= before) & (tpf.time.value <= after)]


def to_lightcurve_with_custom_aperture(tpf, aperture_mask, background_mask):
    """Create a lightcurve from the given TargetPixelFile, with suppplied aperture and background"""
    n_aperture_pixels = aperture_mask.sum()
    aperture_lc = tpf.to_lightcurve(aperture_mask=aperture_mask)  # aperture + background

    n_background_pixels = background_mask.sum()
    background_lc_per_pixel = tpf.to_lightcurve(aperture_mask=background_mask) / n_background_pixels
    background_lc = background_lc_per_pixel * n_aperture_pixels

    corrected_lc = aperture_lc - background_lc
    # OPEN: consider filling in lc metadata, label, etc.

    return corrected_lc, aperture_lc, background_lc


#
# Others
#


def coordinate_like_id_to_coordinate(id, style="decimal"):
    """Given an identifier that describes the target's coordinate, return the cooridnate in string or `SkyCoord` object.
    The type of IDs supported are in the form of '<prefix-of-catalog> J<ra><dec>'.
    Catalogs that adopts such form of ids include PTF1, ASASSN-V, WISE, WISEA, 1SWASP, etc.

    """
    # e.g.,
    # PTF1 J2219+3135 (the shorterned form often found in papers)
    # PTF1 J221910.09+313523.1  (the actual id, with higher precision)
    result = re.findall(r"^.{2,}\s*J(\d{2})(\d{2})([0-9.]*)([+-])(\d{2})(\d{2})([0-9.]*)\s*$", id)
    if len(result) < 1:
        return None
    [ra_hh, ra_mm, ra_ss, dec_sign, dec_deg, dec_min, dec_ss] = result[0]
    coord = SkyCoord(f"{ra_hh} {ra_mm} {ra_ss} {dec_sign} {dec_deg} {dec_min} {dec_ss}", unit=(u.hourangle, u.deg))
    if style is None:
        return coord
    else:
        return coord.to_string(style=style)


#
# Reader for ASAS-SN
#


def read_asas_sn_csv(url=None, asas_sn_uuid=None):
    if url is None:
        if asas_sn_uuid is not None:
            url = f"https://asas-sn.osu.edu/variables/{asas_sn_uuid}/star_data/export?type=csv"
        else:
            raise ValueError("either url or asas_sn_uuid must be supplied")
    else:
        match = re.match(r".*asas-sn.osu.edu/variables/([^/]+)/star_data/export", url)
        if match is not None:
            asas_sn_uuid = match[1]

    tab = Table.read(url, format="ascii")

    # make columns follow Lightkurve convention
    tab.rename_column("mag err", "mag_err")
    tab.rename_column("flux (mJy)", "flux")
    tab.rename_column("flux err", "flux_err")

    tab.sort(keys="hjd")

    tab["mag"].unit = u.mag
    tab["mag_err"].unit = u.mag
    # OPEN: unit of flux is actually mJy. I don't know what it means though
    tab["flux"].unit = u.dimensionless_unscaled
    tab["flux_err"].unit = u.dimensionless_unscaled

    lc = lk.LightCurve(time=Time(tab["hjd"], format="jd", scale="tt"), data=tab)

    lc.meta["FILEURL"] = url
    if asas_sn_uuid is not None:
        lc.meta["ASAS-SN_UUID"] = asas_sn_uuid
        label = f"ASAS-SN {asas_sn_uuid}"
        filters = np.unique(lc.filter)
        if len(filters) > 1:
            label += " , filters: {', '.join(np.unique(lc.filter))}"
        lc.meta["LABEL"] = label

    return lc

<filename>lightkurve_ext.py<gh_stars>1-10
"""
Convenience helpers for `lightkurve` package.
"""

import os
import logging
import math
import json
import re
import warnings
from collections import OrderedDict

from astropy.io import fits
from astropy.coordinates import SkyCoord
from astropy.table import Table
from astropy.time import Time
import astropy.units as u

import numpy as np
from scipy.interpolate import UnivariateSpline

from IPython.display import display, HTML

import lightkurve as lk

import asyncio_compat

log = logging.getLogger(__name__)


def of_sector(lcf_coll, sectorNum):
    res_list = of_sectors(lcf_coll, sectorNum)
    return res_list[0] if len(res_list) > 0 else None


def of_sectors(*args):
    lk_coll_or_sr = args[0]
    if len(args) == 1:
        # when no sectors are specified, return entire collection
        # For convenience: when a notebooks is modified such that
        # a user sometimes use a subset of sectors , and sometimes everything
        # the user can can still use of_sectors() wrapper regardless
        return lk_coll_or_sr
    sector_nums = args[1:]

    if hasattr(lk_coll_or_sr, "sector"):
        return lk_coll_or_sr[np.in1d(lk_coll_or_sr.sector, sector_nums)]
    elif hasattr(lk_coll_or_sr, "table") and lk_coll_or_sr.table["sequence_number"] is not None:
        return lk.SearchResult(lk_coll_or_sr.table[np.in1d(lk_coll_or_sr.table["sequence_number"], sector_nums)])
    else:
        raise TypeError(f"Unsupported type of collection: {type(lk_coll_or_sr)}")


def of_sector_n_around(lk_coll_or_sr, sector_num, num_additions=8):
    def get_sector_for_lk_coll(lk_coll):
        return lk_coll.sector

    def get_sector_for_sr(sr):
        return sr.table["sequence_number"]

    sector_accessor_func = None
    if hasattr(lk_coll_or_sr, "sector"):
        sector_accessor_func = get_sector_for_lk_coll
    elif hasattr(lk_coll_or_sr, "table") and lk_coll_or_sr.table["sequence_number"] is not None:
        sector_accessor_func = get_sector_for_sr
    else:
        raise TypeError(f"Unsupported type of collection: {type(lk_coll_or_sr)}")

    subset_slice = _get_slice_for_of_sector_n_around(
        lk_coll_or_sr,
        sector_accessor_func,
        sector_num,
        num_additions=num_additions,
    )
    if subset_slice is not None:
        return lk_coll_or_sr[subset_slice]
    else:
        return type(lk_coll_or_sr)([])


def _get_slice_for_of_sector_n_around(coll, sector_accessor_func, sector_num, num_additions):
    if sector_num not in sector_accessor_func(coll):
        return None

    idx = np.where(sector_accessor_func(coll) == sector_num)[0][0]
    # if num_additions is odd number, we add one to older sector
    start = max(idx - math.ceil(num_additions / 2), 0)
    end = min(idx + math.floor(num_additions / 2) + 1, len(coll))

    # case the start:end slice does not fill up the requested num_additions,
    # try to fill it up
    cur_slice_size = end - start - 1
    if cur_slice_size < num_additions:
        num_more_needed = num_additions - cur_slice_size
        if start > 0:
            start = max(start - num_more_needed, 0)
        else:
            end = min(end + num_more_needed, len(coll))

    return slice(start, end)


def of_2min_cadences(lcf_coll):
    """Return LightCurveFiles of short, typically 2-minute cadence, only.
    Primary use case is to filter out 20-second files.
    """
    filtered = [lcf for lcf in lcf_coll if "short" == estimate_cadence_type(lcf)]
    return lk.LightCurveCollection(filtered)


def estimate_cadence(lc):
    """Estimate the cadence of a lightcurve by returning the median of a sample"""
    return np.nanmedian(np.diff(lc.time[:100].value))


def map_cadence_type(cadence_in_days):
    long_minimum = 9.9 / 60 / 24  # 10 minutes in days, with some margin of error
    short_minimum = 0.9 / 60 / 24  # 1 minute in days, with some margin of error
    if cadence_in_days is None:
        return None
    if cadence_in_days >= long_minimum:
        return "long"
    if cadence_in_days >= short_minimum:
        return "short"
    return "fast"


def estimate_cadence_type(lc):
    """Estimate the type of cadence to be one of long, short, or fast.
    The definition is the same as ``exptime`` in `lightkurve.search_lightcurve()`.
    """
    return map_cadence_type(estimate_cadence(lc))


def of_tic(lcf_coll, tic):
    """Return LightCurveFiles of the given TIC.

    Useful in case the default MAST result returned nearby targets.
    """
    filtered = [lcf for lcf in lcf_coll if lcf.meta.get("TICID", None) == tic]
    return lk.LightCurveCollection(filtered)


def estimate_object_radius_in_r_jupiter(lc, depth):
    """Return a back of envelope estimate of a companion object's radius."""
    R_JUPITER_IN_R_SUN = 71492 / 695700

    r_star = lc.meta.get("RADIUS")  # assumed to be in R_sun
    if r_star is None or r_star < 0 or depth <= 0:
        return None  # cannot estimate
    r_obj = math.sqrt(r_star * r_star * depth)
    r_obj_in_r_jupiter = r_obj / R_JUPITER_IN_R_SUN
    return r_obj_in_r_jupiter


def download_lightcurves_of_tic_with_priority(tic, download_filter_func=None, download_dir=None):
    """For a given TIC, download lightcurves across all sectors.
    For each sector, download one based on pre-set priority.
    """

    sr_unfiltered = lk.search_lightcurve(f"TIC{tic}", mission="TESS")
    if len(sr_unfiltered) < 1:
        print(f"WARNING: no result found for TIC {tic}")
        return None, None, None

    sr_unfiltered = sr_unfiltered[sr_unfiltered.target_name == str(tic)]  # in case we get some other nearby TICs

    # filter out HLSPs not supported by lightkurve yet
    sr = sr_unfiltered[sr_unfiltered.author != "DIAMANTE"]
    if len(sr) < len(sr_unfiltered):
        print("Note: there are products not supported by Lightkurve, which are excluded from download.")

    # for each sector, filter based on the given priority.
    # - note: prefer QLP over TESS-SPOC because QLP is detrended, with multiple apertures within 1 file
    sr = filter_by_priority(
        sr,
        author_priority=["SPOC", "QLP", "TESS-SPOC"],
        exptime_priority=["short", "long", "fast"],
    )
    num_filtered = len(sr_unfiltered) - len(sr)
    num_fast = len(sr_unfiltered[sr_unfiltered.exptime < 60 * u.second])
    if num_filtered > 0:
        msg = f"{num_filtered} rows filtered"
        if num_fast > 0:
            msg = msg + f" ; {num_fast} fast (20secs) products."
        print(msg)

    display(sr)

    # let caller to optionally further restrict a subset to be downloaded
    sr_to_download = sr
    if download_filter_func is not None:
        sr_to_download = download_filter_func(sr)
        if len(sr_to_download) < len(sr):
            display(
                HTML(
                    """<font style="background-color: yellow;">Note</font>:
SearchResult is further filtered - only a subset will be downloaded."""
                )
            )

    lcf_coll = sr_to_download.download_all(download_dir=download_dir)

    if lcf_coll is not None and len(lcf_coll) > 0:
        print(f"TIC {tic} \t#sectors: {len(lcf_coll)} ; {lcf_coll[0].meta['SECTOR']} - {lcf_coll[-1].meta['SECTOR']}")
        print(
            (
                f"   sector {lcf_coll[-1].meta['SECTOR']}: \t"
                f"camera = {lcf_coll[-1].meta['CAMERA']} ; ccd = {lcf_coll[-1].meta['CCD']}"
            )
        )
    else:
        print(f"TIC {tic}: no data")

    return lcf_coll, sr, sr_unfiltered


def download_lightcurve(
    target,
    mission=("Kepler", "K2", "TESS"),
    exptime="short",
    author="SPOC",
    download_dir=None,
    use_cache="yes",
    display_search_result=True,
):
    """
    Wraps `lightkurve.search_lightcurve()` and the
    subsequent `lightkurve.search.SearchResult.download_all()` calls,
    with the option of caching, so that for a given search,
    if the the result has been downloaded, the cache will be used.

    The parameters all propagate to the underlying `search_lightcurvefile()`
    and `download_all()` calls. The lone exception is `use_cache`.

    Parameters
    ----------
    use_cache : str, must be one of 'yes', or 'no'\n
        OPEN: an option of 'fallback': cache will be used when offline.\n
        OPEN: for now, actual download lightcurve cache will still be used if
        available irrespective of the settings.

    Returns
    -------
    collection : `~lightkurve.collections.Collection` object
        Returns a `~lightkurve.collections.LightCurveCollection`
        containing all lightcurve files that match the criteria
    """

    if use_cache == "no":
        return _search_and_cache(target, mission, exptime, author, download_dir, display_search_result)
    if use_cache == "yes":
        result_file_ids = _load_from_cache_if_any(target, mission, download_dir)
        if result_file_ids is not None:
            result_files = list(map(lambda e: f"{download_dir}/mastDownload/{e}", result_file_ids))
            return lk.collections.LightCurveCollection(list(map(lambda f: lk.read(f), result_files)))
        # else
        return _search_and_cache(target, mission, exptime, author, download_dir, display_search_result)
    # else
    raise ValueError("invalid value for argument use_cache")


# Private helpers for `download_lightcurvefiles`


def _search_and_cache(target, mission, exptime, author, download_dir, display_search_result):
    search_res = lk.search_lightcurve(target=target, mission=mission, exptime=exptime, author=author)
    if len(search_res) < 1:
        return None
    if display_search_result:
        _display_search_result(search_res)
    _cache_search_result_product_identifiers(search_res, download_dir, target, mission)
    return search_res.download_all(quality_bitmask="default", download_dir=download_dir)


def _display_search_result(search_res):
    from IPython.core.display import display

    tab = search_res.table
    # move useful columns to the front
    preferred_cols = ["proposal_id", "target_name", "sequence_number", "t_exptime"]
    colnames_reordered = preferred_cols + [c for c in tab.colnames if c not in preferred_cols]
    display(tab[colnames_reordered])


def _load_from_cache_if_any(target, mission, download_dir):
    key = _get_cache_key(target, mission)
    return _load_search_result_product_identifiers(download_dir, key)


def _cache_search_result_product_identifiers(search_res, download_dir, target, mission):
    key = _get_cache_key(target, mission)
    identifiers = _to_product_identifiers(search_res)
    _save_search_result_product_identifiers(identifiers, download_dir, key)
    return key


def _get_search_result_cache_dir(download_dir):
    # TODO: handle download_dir is None (defaults)
    cache_dir = f"{download_dir}/mastQueries"

    if os.path.isdir(cache_dir):
        return cache_dir

    # else it doesn't exist, make a new cache directory
    try:
        os.mkdir(cache_dir)
    # downloads locally if OS error occurs
    except OSError:
        log.warning(
            "Warning: unable to create {}. "
            "Cache MAST query results to the current "
            "working directory instead.".format(cache_dir)
        )
        cache_dir = "."
    return cache_dir


def _get_cache_key(target, mission):
    # TODO: handle cases the generated key is not a valid filename
    return f"{target}_{mission}_ids"


def _to_product_identifiers(search_res):
    """
    Returns
    -------
    A list of str, constructed from `(obs_collection, obs_id, productFilename)` tuples, that can
    identify cached lightcurve file,s if any.
    """
    return list(
        map(
            lambda e: e["obs_collection"] + "/" + e["obs_id"] + "/" + e["productFilename"],
            search_res.table,
        )
    )


def _save_search_result_product_identifiers(identifiers, download_dir, key):
    resolved_cache_dir = _get_search_result_cache_dir(download_dir)
    filepath = f"{resolved_cache_dir}/{key}.json"
    fp = open(filepath, "w+")
    json.dump(identifiers, fp)
    return filepath


def _load_search_result_product_identifiers(download_dir, key):
    resolved_cache_dir = _get_search_result_cache_dir(download_dir)
    filepath = f"{resolved_cache_dir}/{key}.json"
    try:
        fp = open(filepath, "r")
        return json.load(fp)
    except OSError as err:
        # errno == 2: file not found, typical case of cache miss
        # errno != 2: unexpected error, log a warning
        if err.errno != 2:
            log.warning("Unexpected OSError in retrieving cached search result: {}".format(err))
        return None


def filter_by_priority(
    sr,
    author_priority=["SPOC", "TESS-SPOC", "QLP"],
    exptime_priority=["short", "long", "fast"],
):
    author_sort_keys = {}
    for idx, author in enumerate(author_priority):
        author_sort_keys[author] = idx + 1

    exptime_sort_keys = {}
    for idx, exptime in enumerate(exptime_priority):
        exptime_sort_keys[exptime] = idx + 1

    def calc_filter_priority(row):
        # Overall priority key is in the form of <author_key><exptime_key>, e.g., 101
        # - "01" is the exptime_key
        # - the leading "1" is the author_key, given it is the primary one
        author_default = max(dict(author_sort_keys).values()) + 1
        author_key = author_sort_keys.get(row["author"], author_default) * 100

        # secondary priority
        exptime_default = max(dict(exptime_sort_keys).values()) + 1
        exptime_key = exptime_sort_keys.get(map_cadence_type(row["exptime"] / 60 / 60 / 24), exptime_default)
        return author_key + exptime_key

    sr.table["_filter_priority"] = [calc_filter_priority(r) for r in sr.table]

    # A temporary table that sorts the table by the priority
    sorted_t = sr.table.copy()
    sorted_t.sort(["mission", "_filter_priority"])

    # create an empty table for results, with the same set of columns
    res_t = sr.table[np.zeros(len(sr), dtype=bool)].copy()

    # for each mission (e.g., TESS Sector 01), select a row based on specified priority
    # - select the first row given the table has been sorted by priority
    uniq_missions = list(OrderedDict.fromkeys(sorted_t["mission"]))
    for m in uniq_missions:
        mission_t = sorted_t[sorted_t["mission"] == m]
        # OPEN: if for a given mission, the only row available is not listed in the priorities,
        # the logic still add a row to the result.
        # We might want it to be an option specified by the user.
        res_t.add_row(mission_t[0])

    return lk.SearchResult(table=res_t)


# Download TPF asynchronously


def search_and_download_tpf(*args, **kwargs):
    """Search and Download a TPFs.

    All parameters are passed on ``search_targetpixelfile()``,
    with the exception of ``download_dir`` and ``quality_bitmask``,
    which are passed to ``download_all()`
    """

    # extract download_all() parameters
    download_dir = kwargs.pop("download_dir", None)
    quality_bitmask = kwargs.pop("quality_bitmask", None)
    sr = lk.search_targetpixelfile(*args, **kwargs)  # pass the rest of the argument to search_targetpixelfile
    tpf_coll = sr.download_all(download_dir=download_dir, quality_bitmask=quality_bitmask)
    return tpf_coll, sr


def create_download_tpf_task(*args, **kwargs):
    return asyncio_compat.create_background_task(search_and_download_tpf, *args, **kwargs)


#
# Other misc. extensions
#
def get_bkg_lightcurve(lcf):
    """Returns the background flux, i.e., ``SAP_BKG`` in the file"""
    lc = lcf.copy()
    lc["flux"] = lc["sap_bkg"]
    lc["flux_err"] = lc["sap_bkg_err"]
    lc.label = lc.label + " BKG"
    return lc


def _do_create_quality_issues_mask(quality, flux, flags_included=0b0101001010111111):
    """Returns a boolean array which flags cadences with *issues*.

    The default `flags_included` is a TESS default, based on
    https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags
    """
    if np.issubdtype(quality.dtype, np.integer):
        return np.logical_and(quality & flags_included, np.isfinite(flux))
    else:
        # quality column is not an integer, probably a non-standard product
        return np.zeros_like(quality, dtype=bool)


def create_quality_issues_mask(lc, flags_included=0b0101001010111111):
    """Returns a boolean array which flags cadences with *issues*.

    The default `flags_included` is a TESS default, based on
    https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags
    """

    # use sap_flux when available (it may not be there in some HLSP)
    # we prefer sap_flux over pdcsap_flux as
    # pdcsap_flux is more likely to be NaN (due to exclusion by quality flags)
    if "sap_flux" in lc.colnames:
        flux = lc["sap_flux"]
    else:
        flux = lc.flux

    return _do_create_quality_issues_mask(lc.quality, flux)


def _get_n_truncate_fits_data(lc, before, after, return_columns, return_mask=False):
    with fits.open(lc.filename) as hdu:
        time = hdu[1].data["TIME"]
        mask = (time >= before) & (time < after)
        res = dict()
        for col in return_columns:
            res[col] = hdu[1].data[col][mask]
        if return_mask:
            return res, mask
        else:
            return res


def list_times_w_quality_issues(lc, include_excluded_cadences=False):
    if not include_excluded_cadences:
        mask = create_quality_issues_mask(lc)
        return lc.time[mask], lc.quality[mask]
    else:
        # case we want cadences that have been excluded in the lc object
        # use the underlying fits file

        flux_colname = lc.meta.get("FLUX_ORIGIN", "sap_flux")
        if flux_colname == "pdcsap_flux":
            flux_colname = "sap_flux"  # pdcsap_flux would not have values in excluded cadences, defeating the purpose

        # data is truncated if the given lc is truncated
        # TODO: cadences excluded before lc.time.min() or after lc.time.max() will still be missing.
        data = _get_n_truncate_fits_data(lc, lc.time.min().value, lc.time.max().value, ["time", "quality", flux_colname])
        mask = _do_create_quality_issues_mask(data["quality"], data[flux_colname])
        return data["time"][mask], data["quality"][mask]


def list_transit_times(t0, period, steps_or_num_transits=range(0, 10), return_string=False):
    """List the transit times based on the supplied transit parameters"""
    if isinstance(steps_or_num_transits, int):
        steps = range(0, steps_or_num_transits)
    else:
        steps = steps_or_num_transits
    times = [t0 + period * i for i in steps]
    if return_string:
        return ",".join(map(str, times))
    else:
        return times


def get_segment_times_idx(times, break_tolerance=5):
    """Segment the input array of times into segments due to data gaps. Return the indices of the segments.

    The minimal gap size is determined by `break_tolerance`.

    The logic is adapted from `LightCurve.flatten`
    """
    if hasattr(times, "value"):  # convert astropy Time to raw values if needed
        times = times.value
    dt = times[1:] - times[0:-1]
    with warnings.catch_warnings():  # Ignore warnings due to NaNs
        warnings.simplefilter("ignore", RuntimeWarning)
        cut = np.where(dt > break_tolerance * np.nanmedian(dt))[0] + 1
    low = np.append([0], cut)
    high = np.append(cut, len(times))
    return (low, high)


def get_segment_times(times, **kwargs):
    if hasattr(times, "value"):  # convert astropy Time to raw values if needed
        times = times.value
    low, high = get_segment_times_idx(times, **kwargs)
    # add a small 1e-10 to end so that the end time is exclusive (follow convention in range)
    return [(times[lo], times[hi - 1] + 1e-10) for lo, hi in zip(low, high)]


def get_transit_times_in_range(t0, period, start, end):
    t_start = t0 + math.ceil((start - t0) / period) * period
    num_t = math.ceil((end - t_start) / period)
    return [t_start + period * i for i in range(num_t)]


def get_transit_times_in_lc(lc, t0, period, return_string=False, **kwargs):
    """Get the transit times with observations of the given lightcurve, based on the supplied transit parameters.

    The method will exclude the times where there is no observation due to data gaps.
    """

    lc = lc.remove_nans()  # exclude cadences with no flux.
    # break up the times to exclude times in gap
    times_list = get_segment_times(lc.time, **kwargs)
    transit_times = []
    for start, end in times_list:
        transit_times.extend(get_transit_times_in_range(t0, period, start, end))
    if return_string:
        return ",".join(map(str, transit_times))
    else:
        return transit_times


def stitch(lcf_coll, **kwargs):
    """Wrapper over native stitch(), and tweak the metadata so that it behaves like a typical single-sector lightcurve."""

    def update_meta_if_exists_in(lc_src, keys):
        for key in keys:
            val = lc_src.meta.get(key, None)
            if val is not None:
                lc_stitched.meta[key] = val

    def safe_del_meta(key):
        if lc_stitched.meta.get(key, None) is not None:
            del lc_stitched.meta[key]

    lc_stitched = lcf_coll.stitch(**kwargs)

    # now update the metadata

    lc_stitched.meta["STITCHED"] = True

    # update observation start/stop dates
    update_meta_if_exists_in(lcf_coll[0], ("TSTART", "DATE-OBS"))
    update_meta_if_exists_in(lcf_coll[-1], ("TSTOP", "DATE-END"))

    # TODO: recalculate TELAPSE, LIVETIME, DEADC (which is LIVETIME / TELAPSE)

    safe_del_meta("FILENAME")  # don't associate it with a file anymore

    # record the sectors stitched and the associated metadata
    sector_list = [lc.meta.get("SECTOR") for lc in lcf_coll if lc.meta.get("SECTOR") is not None]
    meta_list = [lc.meta for lc in lcf_coll if lc.meta.get("SECTOR") is not None]
    if len(sector_list) > 0:
        lc_stitched.meta["SECTORS"] = sector_list
        meta_dict = dict()
        for sector, meta in zip(sector_list, meta_list):
            meta_dict[sector] = meta.copy()
        lc_stitched.meta["HEADERS_ORIGINAL"] = meta_dict

    return lc_stitched


def to_window_length_for_2min_cadence(length_day):
    """Helper for LightCurve.flatten().
    Return a `window_length` for the given number of days, assuming the data has 2-minute cadence."""
    res = math.floor(720 * length_day)
    if res % 2 == 0:
        res += 1  # savgol_filter window length must be odd number
    return res


# detrend using spline
# Based on:  https://github.com/barentsen/kepler-athenaeum-tutorial/blob/master/how-to-find-a-planet-tutorial.ipynb
def flatten_with_spline_normalized(lc, return_trend=False, **kwargs):
    lc = lc.remove_nans()
    spline = UnivariateSpline(lc.time, lc.flux, **kwargs)
    trend = spline(lc.time)
    # detrended = lc.flux - trend
    detrended_relative = 100 * ((lc.flux / trend) - 1) + 100  # in percentage
    lc_flattened = lc.copy()
    lc_flattened.flux = detrended_relative
    lc_flattened.flux_unit = "percent"
    if not return_trend:
        return lc_flattened
    else:
        lc_trend = lc.copy()
        lc_trend.flux = trend
        return (lc_flattened, lc_trend)


def _lksl_statistics(ts):
    """Compute LKSL Statistics of the given (time-series) values.
    Useful to compare the noises in a folded lightcurve.

    Based on https://arxiv.org/pdf/1901.00009.pdf , equation 4.
    See https://www.aanda.org/articles/aa/pdf/2002/17/aa2208.pdf for more information.
    """
    ts = ts[~np.isnan(ts)]

    vector_length_sq_sum = np.square(np.diff(ts)).sum()
    if len(ts) > 2:
        # to fully utilize the data by including the vector length between the last and first measurement
        # section 2 of Clarke, 2002
        vector_length_sq_sum += np.square(ts[-1] - ts[0])

    diff_from_mean_sq_sum = np.square(ts - np.mean(ts)).sum()
    if diff_from_mean_sq_sum == 0:  # to avoid boundary case that'd cause division by zero
        diff_from_mean_sq_sum = 1e-10

    return (vector_length_sq_sum / diff_from_mean_sq_sum) * (len(ts) - 1) / (2 * len(ts))


def lksl_statistics(lc, column="flux"):
    return _lksl_statistics(lc[column].value)


def tess_flux_to_mag(flux):
    """Convert flux from TESS observation to magnitude."""
    # Based on https://heasarc.gsfc.nasa.gov/docs/tess/observing-technical.html#saturation
    # TESS CCDs produce 15000 e/s for magnitude 10 light source

    if isinstance(flux, u.Quantity):
        flux_raw = flux.to(u.electron / u.second).value
    else:
        flux_raw = flux

    # np.log10 does not work on Quantity, unless it's dimensionless
    res_raw = 10 - 2.5 * np.log10((flux_raw / 15000))

    if isinstance(flux, u.Quantity):
        return res_raw * u.mag
    else:
        return res_raw


def mag_to_tess_flux(mag):
    """Convert magnitude to flux in TESS observation."""
    if isinstance(mag, u.Quantity):
        mag_raw = mag.to(u.mag).value
    else:
        mag_raw = mag

    flux_raw = (10 ** ((10 - mag_raw) / 2.5)) * 15000
    if isinstance(mag, u.Quantity):
        return flux_raw * ((u.electron / u.second))
    else:
        return flux_raw


def select_flux(lc, flux_cols):
    """Return a Lightcurve object with the named column as the flux column.

    flux_cols: either a column name (string), or a list of prioritized column names
    such that the first one that the lightcurve contains will be used.
    """

    def _to_lc_with_1flux(lc, flux_1col):
        flux_1col = flux_1col.lower()
        if "flux" == flux_1col:
            return lc
        elif flux_1col in lc.colnames:
            return lc.select_flux(flux_1col)
        else:
            return None

    if isinstance(flux_cols, str):
        flux_cols = [flux_cols]

    for flux_1col in flux_cols:
        res = _to_lc_with_1flux(lc, flux_1col)
        if res is not None:
            return res
    raise ValueError(f"'column {flux_cols}' not found")


#
# TargetPixelFile helpers
#


def truncate(tpf: lk.targetpixelfile.TargetPixelFile, before: float, after: float) -> lk.targetpixelfile.TargetPixelFile:
    return tpf[(tpf.time.value >= before) & (tpf.time.value <= after)]


def to_lightcurve_with_custom_aperture(tpf, aperture_mask, background_mask):
    """Create a lightcurve from the given TargetPixelFile, with suppplied aperture and background"""
    n_aperture_pixels = aperture_mask.sum()
    aperture_lc = tpf.to_lightcurve(aperture_mask=aperture_mask)  # aperture + background

    n_background_pixels = background_mask.sum()
    background_lc_per_pixel = tpf.to_lightcurve(aperture_mask=background_mask) / n_background_pixels
    background_lc = background_lc_per_pixel * n_aperture_pixels

    corrected_lc = aperture_lc - background_lc
    # OPEN: consider filling in lc metadata, label, etc.

    return corrected_lc, aperture_lc, background_lc


#
# Others
#


def coordinate_like_id_to_coordinate(id, style="decimal"):
    """Given an identifier that describes the target's coordinate, return the cooridnate in string or `SkyCoord` object.
    The type of IDs supported are in the form of '<prefix-of-catalog> J<ra><dec>'.
    Catalogs that adopts such form of ids include PTF1, ASASSN-V, WISE, WISEA, 1SWASP, etc.

    """
    # e.g.,
    # PTF1 J2219+3135 (the shorterned form often found in papers)
    # PTF1 J221910.09+313523.1  (the actual id, with higher precision)
    result = re.findall(r"^.{2,}\s*J(\d{2})(\d{2})([0-9.]*)([+-])(\d{2})(\d{2})([0-9.]*)\s*$", id)
    if len(result) < 1:
        return None
    [ra_hh, ra_mm, ra_ss, dec_sign, dec_deg, dec_min, dec_ss] = result[0]
    coord = SkyCoord(f"{ra_hh} {ra_mm} {ra_ss} {dec_sign} {dec_deg} {dec_min} {dec_ss}", unit=(u.hourangle, u.deg))
    if style is None:
        return coord
    else:
        return coord.to_string(style=style)


#
# Reader for ASAS-SN
#


def read_asas_sn_csv(url=None, asas_sn_uuid=None):
    if url is None:
        if asas_sn_uuid is not None:
            url = f"https://asas-sn.osu.edu/variables/{asas_sn_uuid}/star_data/export?type=csv"
        else:
            raise ValueError("either url or asas_sn_uuid must be supplied")
    else:
        match = re.match(r".*asas-sn.osu.edu/variables/([^/]+)/star_data/export", url)
        if match is not None:
            asas_sn_uuid = match[1]

    tab = Table.read(url, format="ascii")

    # make columns follow Lightkurve convention
    tab.rename_column("mag err", "mag_err")
    tab.rename_column("flux (mJy)", "flux")
    tab.rename_column("flux err", "flux_err")

    tab.sort(keys="hjd")

    tab["mag"].unit = u.mag
    tab["mag_err"].unit = u.mag
    # OPEN: unit of flux is actually mJy. I don't know what it means though
    tab["flux"].unit = u.dimensionless_unscaled
    tab["flux_err"].unit = u.dimensionless_unscaled

    lc = lk.LightCurve(time=Time(tab["hjd"], format="jd", scale="tt"), data=tab)

    lc.meta["FILEURL"] = url
    if asas_sn_uuid is not None:
        lc.meta["ASAS-SN_UUID"] = asas_sn_uuid
        label = f"ASAS-SN {asas_sn_uuid}"
        filters = np.unique(lc.filter)
        if len(filters) > 1:
            label += " , filters: {', '.join(np.unique(lc.filter))}"
        lc.meta["LABEL"] = label

    return lc

Convenience helpers for `lightkurve` package. # when no sectors are specified, return entire collection # For convenience: when a notebooks is modified such that # a user sometimes use a subset of sectors , and sometimes everything # the user can can still use of_sectors() wrapper regardless # if num_additions is odd number, we add one to older sector # case the start:end slice does not fill up the requested num_additions, # try to fill it up Return LightCurveFiles of short, typically 2-minute cadence, only.     Primary use case is to filter out 20-second files. Estimate the cadence of a lightcurve by returning the median of a sample # 10 minutes in days, with some margin of error # 1 minute in days, with some margin of error Estimate the type of cadence to be one of long, short, or fast.     The definition is the same as ``exptime`` in `lightkurve.search_lightcurve()`. Return LightCurveFiles of the given TIC.      Useful in case the default MAST result returned nearby targets. Return a back of envelope estimate of a companion object's radius. # assumed to be in R_sun # cannot estimate For a given TIC, download lightcurves across all sectors.     For each sector, download one based on pre-set priority. # in case we get some other nearby TICs # filter out HLSPs not supported by lightkurve yet # for each sector, filter based on the given priority. # - note: prefer QLP over TESS-SPOC because QLP is detrended, with multiple apertures within 1 file # let caller to optionally further restrict a subset to be downloaded <font style="background-color: yellow;">Note</font>: SearchResult is further filtered - only a subset will be downloaded. #sectors: {len(lcf_coll)} ; {lcf_coll[0].meta['SECTOR']} - {lcf_coll[-1].meta['SECTOR']}") Wraps `lightkurve.search_lightcurve()` and the     subsequent `lightkurve.search.SearchResult.download_all()` calls,     with the option of caching, so that for a given search,     if the the result has been downloaded, the cache will be used.      The parameters all propagate to the underlying `search_lightcurvefile()`     and `download_all()` calls. The lone exception is `use_cache`.      Parameters     ----------     use_cache : str, must be one of 'yes', or 'no'\n         OPEN: an option of 'fallback': cache will be used when offline.\n         OPEN: for now, actual download lightcurve cache will still be used if         available irrespective of the settings.      Returns     -------     collection : `~lightkurve.collections.Collection` object         Returns a `~lightkurve.collections.LightCurveCollection`         containing all lightcurve files that match the criteria # else # else # Private helpers for `download_lightcurvefiles` # move useful columns to the front # TODO: handle download_dir is None (defaults) # else it doesn't exist, make a new cache directory # downloads locally if OS error occurs # TODO: handle cases the generated key is not a valid filename Returns     -------     A list of str, constructed from `(obs_collection, obs_id, productFilename)` tuples, that can     identify cached lightcurve file,s if any. # errno == 2: file not found, typical case of cache miss # errno != 2: unexpected error, log a warning # Overall priority key is in the form of <author_key><exptime_key>, e.g., 101 # - "01" is the exptime_key # - the leading "1" is the author_key, given it is the primary one # secondary priority # A temporary table that sorts the table by the priority # create an empty table for results, with the same set of columns # for each mission (e.g., TESS Sector 01), select a row based on specified priority # - select the first row given the table has been sorted by priority # OPEN: if for a given mission, the only row available is not listed in the priorities, # the logic still add a row to the result. # We might want it to be an option specified by the user. # Download TPF asynchronously Search and Download a TPFs.      All parameters are passed on ``search_targetpixelfile()``,     with the exception of ``download_dir`` and ``quality_bitmask``,     which are passed to ``download_all()` # extract download_all() parameters # pass the rest of the argument to search_targetpixelfile # # Other misc. extensions # Returns the background flux, i.e., ``SAP_BKG`` in the file Returns a boolean array which flags cadences with *issues*.      The default `flags_included` is a TESS default, based on     https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags # quality column is not an integer, probably a non-standard product Returns a boolean array which flags cadences with *issues*.      The default `flags_included` is a TESS default, based on     https://outerspace.stsci.edu/display/TESS/2.0+-+Data+Product+Overview#id-2.0DataProductOverview-Table:CadenceQualityFlags # use sap_flux when available (it may not be there in some HLSP) # we prefer sap_flux over pdcsap_flux as # pdcsap_flux is more likely to be NaN (due to exclusion by quality flags) # case we want cadences that have been excluded in the lc object # use the underlying fits file # pdcsap_flux would not have values in excluded cadences, defeating the purpose # data is truncated if the given lc is truncated # TODO: cadences excluded before lc.time.min() or after lc.time.max() will still be missing. List the transit times based on the supplied transit parameters Segment the input array of times into segments due to data gaps. Return the indices of the segments.      The minimal gap size is determined by `break_tolerance`.      The logic is adapted from `LightCurve.flatten` # convert astropy Time to raw values if needed # Ignore warnings due to NaNs # convert astropy Time to raw values if needed # add a small 1e-10 to end so that the end time is exclusive (follow convention in range) Get the transit times with observations of the given lightcurve, based on the supplied transit parameters.      The method will exclude the times where there is no observation due to data gaps. # exclude cadences with no flux. # break up the times to exclude times in gap Wrapper over native stitch(), and tweak the metadata so that it behaves like a typical single-sector lightcurve. # now update the metadata # update observation start/stop dates # TODO: recalculate TELAPSE, LIVETIME, DEADC (which is LIVETIME / TELAPSE) # don't associate it with a file anymore # record the sectors stitched and the associated metadata Helper for LightCurve.flatten().     Return a `window_length` for the given number of days, assuming the data has 2-minute cadence. # savgol_filter window length must be odd number # detrend using spline # Based on:  https://github.com/barentsen/kepler-athenaeum-tutorial/blob/master/how-to-find-a-planet-tutorial.ipynb # detrended = lc.flux - trend # in percentage Compute LKSL Statistics of the given (time-series) values.     Useful to compare the noises in a folded lightcurve.      Based on https://arxiv.org/pdf/1901.00009.pdf , equation 4.     See https://www.aanda.org/articles/aa/pdf/2002/17/aa2208.pdf for more information. # to fully utilize the data by including the vector length between the last and first measurement # section 2 of Clarke, 2002 # to avoid boundary case that'd cause division by zero Convert flux from TESS observation to magnitude. # Based on https://heasarc.gsfc.nasa.gov/docs/tess/observing-technical.html#saturation # TESS CCDs produce 15000 e/s for magnitude 10 light source # np.log10 does not work on Quantity, unless it's dimensionless Convert magnitude to flux in TESS observation. Return a Lightcurve object with the named column as the flux column.      flux_cols: either a column name (string), or a list of prioritized column names     such that the first one that the lightcurve contains will be used. # # TargetPixelFile helpers # Create a lightcurve from the given TargetPixelFile, with suppplied aperture and background # aperture + background # OPEN: consider filling in lc metadata, label, etc. # # Others # Given an identifier that describes the target's coordinate, return the cooridnate in string or `SkyCoord` object.     The type of IDs supported are in the form of '<prefix-of-catalog> J<ra><dec>'.     Catalogs that adopts such form of ids include PTF1, ASASSN-V, WISE, WISEA, 1SWASP, etc. # e.g., # PTF1 J2219+3135 (the shorterned form often found in papers) # PTF1 J221910.09+313523.1  (the actual id, with higher precision) # # Reader for ASAS-SN # # make columns follow Lightkurve convention # OPEN: unit of flux is actually mJy. I don't know what it means though