utils.databases.lightcurves

  1"""
  2Getting light curves data.
  3"""
  4
  5import lightkurve
  6from astropy.table import Table
  7import pandas
  8from pandera import pandas as pandera
  9import numpy
 10import pathlib
 11import re
 12from packaging.version import Version
 13
 14from typing import Optional, Dict, List, Pattern, Literal
 15
 16from ..files import file as fl
 17from ..logs.log import logger
 18
 19# apparently, one cannot set long/short threshold,
 20# hence this dictionary
 21#
 22# there are actually more authors available,
 23# but we are only interested in these
 24#
 25authors: Dict[str, Dict] = {
 26    "Kepler":
 27    {
 28        "mission": "Kepler",
 29        "cadence":
 30        {
 31            "long": [1800],
 32            "short": [60]
 33        }
 34    },
 35    "K2":
 36    {
 37        "mission": "K2",
 38        "cadence":
 39        {
 40            "long": [1800],
 41            "short": [60]
 42        }
 43    },
 44    "SPOC":
 45    {
 46        "mission": "TESS",
 47        "cadence":
 48        {
 49            "long": [600],
 50            "short": [120],
 51            "fast": [20]
 52        }
 53    },
 54    "TESS-SPOC":
 55    {
 56        "mission": "TESS",
 57        "cadence":
 58        {
 59            "long": []  # any cadence is long
 60        }
 61    }
 62}
 63"""
 64Dictionary of authors, their cadence values and mapping to missions.
 65"""
 66
 67missionSectorRegExes: Dict[str, Pattern] = {
 68    "Kepler": re.compile(
 69        r"^Kepler\s\w+\s(\d+)$"  # Kepler Quarter 15
 70    ),
 71    "K2": re.compile(
 72        r"^K2\s\w+\s(\d+)$"  # K2 Campaign 12
 73    ),
 74    "TESS": re.compile(
 75        r"^TESS\s\w+\s(\d+)$"  # TESS Sector 40
 76    )
 77}
 78"""
 79Dictionary of regular expressions for extracting sectors.
 80"""
 81
 82lightCurveFluxTableSchema = pandera.DataFrameSchema(
 83    {
 84        "time": pandera.Column(numpy.float64),
 85        "flux": pandera.Column(numpy.float32, nullable=True),
 86        "fluxError": pandera.Column(numpy.float32, nullable=True)
 87    },
 88    index=pandera.Index(int, unique=True),
 89    strict=True,  # only specified columns are allowed
 90    coerce=False  # do not cast other types to the specified one
 91)
 92"""
 93Table schema for light curve fluxes.
 94"""
 95
 96
 97def getLightCurveStats(
 98    starName: str,
 99    detailed: bool = True
100) -> Dict[str, Dict]:
101    """
102    Gather statistics about available cadence values for a given star.
103
104    If `detailed` is set to `False`, then function will skip collecting
105    cadence values count by sectors, so resulting statistics will only
106    contain total count of values.
107
108    Example:
109
110    ``` py
111    from phab.utils.databases import lightcurves
112
113    stats = lightcurves.getLightCurveStats("Kepler-114")
114    if not stats:
115        print("Didn't find any results for this star")
116    else:
117        missionName = "Kepler"
118        cadenceType = "long"
119        sectors = stats.get(
120            missionName,
121            {}
122        ).get(cadenceType)
123        if sectors is None:
124            print(
125                " ".join((
126                    "There doesn't seem to be any sectors",
127                    f"with [{cadenceType}] cadence by [{missionName}]"
128                ))
129            )
130        else:
131            totalProperty = "total"
132            sectorsCount = sectors.get(totalProperty)
133            if sectorsCount is None:
134                print(
135                    " ".join((
136                        f"For some reason, the [{totalProperty}] property",
137                        f"is missing from the [{cadenceType}] cadence",
138                        f"collection by [{missionName}]"
139                    ))
140                )
141            else:
142                print(
143                    " ".join((
144                        f"Total amount of sectors with [{cadenceType}]",
145                        f"cadence by [{missionName}]: {sectorsCount}",
146                    ))
147                )
148                bySectors = sectors.get("by-sectors")
149                if bySectors is None:
150                    print(
151                        " ".join((
152                            "For some reason, the [total] property is missing",
153                            f"from the [{cadenceType}] cadence collection",
154                            f"by [{missionName}]"
155                        ))
156                    )
157                else:
158                    for s in bySectors:
159                        print(f"- {s}: {bySectors[s]}")
160    ```
161    """
162    stats: Dict[str, Dict] = {}
163
164    lghtcrvs = lightkurve.search_lightcurve(
165        starName,
166        author=tuple(authors.keys())
167    )
168    if len(lghtcrvs) != 0:
169        tbl: pandas.DataFrame = lghtcrvs.table.to_pandas()[
170            ["author", "exptime", "mission"]
171        ]
172        logger.debug(tbl)
173
174        author: str  # for mypy, but even then it is not happy with something else
175        for author, group in (tbl.groupby("author")):  # type:ignore[assignment] # ya hz
176            if author not in authors:
177                raise ValueError(f"Unknown author: {author}")
178            mission = authors[author]["mission"]
179            if not stats.get(mission):
180                stats[mission] = {}
181            for cadence in ["long", "short", "fast"]:
182                if cadence in authors[author]["cadence"]:
183                    stats[mission][cadence] = {}
184                    cadenceValues: List[int] = (
185                        authors[author]["cadence"][cadence]
186                    )
187                    cadences: pandas.DataFrame
188                    if len(cadenceValues) > 0:  # take only specified values
189                        # perhaps both of these should be normalized to int
190                        cadences = group.query("exptime == @cadenceValues")
191                    else:  # any value is good
192                        cadences = group
193
194                    # total count
195                    stats[mission][cadence]["total"] = len(cadences)
196
197                    if detailed:
198                        # count by sectors
199                        stats[mission][cadence]["by-sectors"] = {}
200                        for m in cadences["mission"]:
201                            # logger.debug(cadences.query("mission == @m")[
202                            #     "exptime"
203                            # ].values)
204                            sectorMatch = re.search(
205                                missionSectorRegExes[mission],
206                                m
207                            )
208                            if not sectorMatch:
209                                raise ValueError(
210                                    " ".join((
211                                        "Couldn't extract sector from",
212                                        f"this mission value: {m}"
213                                    ))
214                                )
215                            sector = sectorMatch.group(1)
216                            if not stats[mission][cadence]["by-sectors"].get(
217                                sector
218                            ):  # this sector hasn't been added yet
219                                stats[mission][cadence]["by-sectors"][
220                                    sector
221                                ] = {}
222                                # save the cadence/exptime too (assuming
223                                # that it is the same for every sector entry)
224                                stats[mission][cadence]["by-sectors"][sector][
225                                    "exptime"
226                                ] = cadences.query("mission == @m")[
227                                    "exptime"
228                                ].values[0]  # there must be a better way
229                            try:
230                                stats[mission][cadence][
231                                    "by-sectors"
232                                ][sector]["count"] += 1
233                            except KeyError:
234                                stats[mission][cadence][
235                                    "by-sectors"
236                                ][sector]["count"] = 1
237    return stats
238
239
240def getLightCurveIDs(
241    starName: str
242) -> Dict[str, List[str]]:
243    """
244    Based on available cadence values statistics for a given star,
245    get names of missions and cadences. For instance, in order to pass
246    them to `altaipony.lcio.from_mast()`.
247
248    Example:
249
250    ``` py
251    from phab.utils.databases import lightcurves
252    from altaipony.lcio import from_mast
253
254    starName = "LTT 1445 A"
255    lightCurveIDs = {}
256
257    try:
258        lightCurveIDs = lightcurves.getLightCurveIDs(starName)
259    except ValueError as ex:
260        print(f"Failed to get light curves missons and cadences. {ex}")
261        raise
262    if not lightCurveIDs:
263        raise ValueError("Didn't find any results for this star")
264    #print(lightCurveIDs)
265
266    for m in lightCurveIDs.keys():
267        #print(f"Mission: {m}")
268        for c in lightCurveIDs[m]:
269            #print(f"- {c}")
270            flc = from_mast(
271                starName,
272                mode="LC",
273                cadence=c,
274                mission=m
275            )
276            #print(flc)
277    ```
278    """
279    lightCurveIDs: Dict[str, List[str]] = {}
280
281    stats: Dict[str, Dict] = getLightCurveStats(
282        starName,
283        detailed=False
284    )
285    if not stats:
286        raise ValueError("Didn't find any results for this star")
287
288    # the order matters, it goes from most important to least important,
289    # and in fact long cadence is so not important that it is discarded
290    # if there is fast or short cadence available
291    cadencePriority = ["fast", "short", "long"]
292
293    for m in stats.keys():
294        lightCurveIDs[m] = []
295        priorityThreshold = 0
296        for cp in cadencePriority:
297            # if there is already fast or short cadence in the list,
298            # don't take long cadence (except for mission K2, because
299            # its long cadence is what's most important even if
300            # there are also fast and short ones)
301            if any(lightCurveIDs[m]) and priorityThreshold > 1 and m != "K2":
302                break
303            if cp in stats[m]:
304                # print(f"Count [{cp}]: {stats[m][cp]['total']}")
305                totalCnt = stats[m][cp].get("total")
306                if totalCnt and totalCnt != 0:
307                    lightCurveIDs[m].append(cp)
308                # else:
309                #     print(
310                #         " ".join((
311                #             f"[WARNING] The [{cp}] cadence count",
312                #             f"in [{m}] is 0 (or missing)"
313                #         ))
314                #     )
315            priorityThreshold += 1
316
317    return lightCurveIDs
318
319
320def fitsToPandas(
321    fitsFilePath: str,
322    fitsType: Optional[Literal["tess", "kepler"]] = None,
323    qualityBitmask: Literal["none", "default", "hard", "hardest"] = "default",
324    dropNanTimes: bool = True,
325    convertTimesToSeconds: bool = False
326) -> pandas.DataFrame:
327    """
328    Open a generic light curves [FITS](https://en.wikipedia.org/wiki/FITS) file
329    and create a Pandas table from it. Only the fluxes, their times
330    and errors columns are taken.
331
332    Handles the big/little endians problem when converting from FITS to Pandas.
333
334    Example:
335
336    ``` py
337    from phab.utils.databases import lightcurves
338
339    pnd = lightcurves.fitsToPandas(
340        "./data/tess2019006130736-s0007-0000000266744225-0131-s_lc.fits",
341        fitsType="tess",
342        qualityBitmask="default",
343        dropNanTimes=True,
344        convertTimesToSeconds=True
345    )
346
347    #print(pnd)
348    ```
349    """
350    lc = None
351    fitsFile: Optional[pathlib.Path] = fl.fileExists(fitsFilePath)
352    if fitsFile is None:
353        raise ValueError(
354            f"Provided path to [{fitsFilePath}] seems to be wrong"
355        )
356    else:
357        lc = Table.read(fitsFile)
358
359    # exclude values which do not satisfy the required quality
360    if fitsType is not None:
361        msk = None
362        if fitsType == "tess":
363            msk = lightkurve.utils.TessQualityFlags.create_quality_mask(
364                quality_array=lc["QUALITY"],
365                bitmask=qualityBitmask
366            )
367        elif fitsType == "kepler":
368            msk = lightkurve.utils.KeplerQualityFlags.create_quality_mask(
369                quality_array=lc["QUALITY"],
370                bitmask=qualityBitmask
371            )
372        else:
373            print(
374                " ".join((
375                    "[WARNING] Unknown FITS type, don't know",
376                    "which quality mask to use"
377                ))
378            )
379        lc = lc[msk]
380
381    narr = numpy.array(lc)
382    # FITS stores data in big-endian, but pandas works with little-endian,
383    # so the byte order needs to be swapped
384    # https://stackoverflow.com/a/30284033/1688203
385    if Version(numpy.__version__) > Version("1.26.4"):
386        # if that doesn't work, then you might need to downgrade to 1.26.4
387        narr = narr.view(narr.dtype.newbyteorder()).byteswap()
388    else:
389        narr = narr.byteswap().newbyteorder()
390
391    # astropy.time does not(?) support NaN
392    if dropNanTimes:
393        nantimes = numpy.isnan(narr["TIME"].data)
394        if numpy.any(nantimes):
395            print(
396                " ".join((
397                    f"[DEBUG] {numpy.sum(nantimes)} rows were excluded,",
398                    "because their time values are NaN"
399                ))
400            )
401        narr = narr[~nantimes]
402
403    # apparently, one cannot just take columns from `lc`/`narr` directly,
404    # hence this intermediate table
405    pndraw = pandas.DataFrame(narr)
406    logger.debug(f"Light curve table columns: {pndraw.columns}")
407
408    flux = pandas.DataFrame(
409        columns=[
410            "time",
411            "flux",
412            "fluxError"
413        ]
414    )
415    flux["time"] = pndraw["TIME"]
416    flux["flux"] = pndraw["PDCSAP_FLUX"]
417    flux["fluxError"] = pndraw["PDCSAP_FLUX_ERR"]
418
419    # in case excluding NaN times right after `Table.read()` is less efficient
420    # if dropNanTimes:
421    #     flux = flux.dropna(subset=["time"])
422
423    if convertTimesToSeconds:
424        flux["time"] = flux["time"] * 24 * 60 * 60
425
426    if dropNanTimes:
427        lightCurveFluxTableSchema.validate(flux)
428    else:
429        lightCurveFluxTableSchema.update_column(
430            "time",
431            dtype=numpy.float64,
432            nullable=True
433        ).validate(flux)
434
435    return flux
436
437
438def lightCurveTessToPandas(
439    lightKurve: lightkurve.lightcurve.TessLightCurve,
440    convertTimesToSeconds: bool = False
441) -> pandas.DataFrame:
442    """
443    Converting a TESS light curve object to a Pandas table. In general,
444    it does almost the same thing as
445    `utils.databases.lightcurves.fitsToPandas()`,
446    but here there it uses a TESS-specific reading function, and also
447    there is no need to drop NaN times "manually" (*and fiddle with endians?*).
448
449    Example:
450
451    ``` py
452    from phab.utils.databases import lightcurves
453    import lightkurve
454
455    downloadLC: bool = False
456    lc = None
457    if downloadLC:
458        search_result = lightkurve.search_lightcurve(
459            "Karmn J07446+035",
460            author="SPOC",
461            cadence="short"
462        )
463        lc = search_result[0].download(
464            quality_bitmask="default"
465        )
466    else:
467        lc = lightkurve.TessLightCurve.read(
468            "./data/tess2019006130736-s0007-0000000266744225-0131-s_lc.fits",
469            quality_bitmask="default"
470        )
471
472    pnd = lightcurves.lightCurveTessToPandas(lc, convertTimesToSeconds=True)
473
474    #print(pnd)
475    ```
476    """
477    pndraw = lightKurve.to_pandas()
478    logger.debug(f"Light curve table columns: {pndraw.columns}")
479
480    flux = pandas.DataFrame(
481        columns=[
482            "time",
483            "flux",
484            "fluxError"
485        ]
486    )
487
488    flux["time"] = pndraw.index
489    flux["flux"] = pndraw["pdcsap_flux"].values
490    flux["fluxError"] = pndraw["pdcsap_flux_err"].values
491
492    if convertTimesToSeconds:
493        flux["time"] = flux["time"] * 24 * 60 * 60
494
495    lightCurveFluxTableSchema.validate(flux)
496
497    return flux