tree wide: way too many changes, im sorry

main/util/MinimalLightCurve.py

@@ -17,6 +17,10 @@ from astropy.time.formats import TimeFromEpoch
 
 log = logging.getLogger(__name__)
 
+class LightkurveWarning(Warning):
+    """Class for all Lightkurve warnings."""
+    pass
+
 class TimeBKJD(TimeFromEpoch):
     name = 'bkjd'
     unit = 1.0
@@ -402,6 +406,14 @@ class TessQualityFlags(QualityFlags):
         32768: "Insufficient Targets for Error Correction Exclude",
     }
 
+def validate_method(method, supported_methods):
+    method = method.lower()
+    if method in supported_methods:
+        return method
+    raise ValueError(
+        "method '{}' is not supported; "
+        "must be one of {}".format(method, supported_methods)
+    )
 
 def _is_dict_like(data1):
     return hasattr(data1, "keys") and callable(getattr(data1, "keys"))
@@ -879,6 +891,195 @@ class LightCurve(TimeSeries):
             return flatten_lc, trend_lc
         return flatten_lc
 
+    def to_periodogram(self, method="lombscargle", **kwargs):
+        supported_methods = ["ls", "bls", "lombscargle", "boxleastsquares"]
+        method = validate_method(method.replace(" ", ""), supported_methods)
+        if method in ["bls", "boxleastsquares"]:
+            from .MinimalPeriodogram import BoxLeastSquaresPeriodogram
+
+            return BoxLeastSquaresPeriodogram.from_lightcurve(lc=self, **kwargs)
+        else:
+            from .MinimalPeriodogram import LombScarglePeriodogram
+
+            return LombScarglePeriodogram.from_lightcurve(lc=self, **kwargs)
+
+    def fold(
+        self,
+        period=None,
+        epoch_time=None,
+        epoch_phase=0,
+        wrap_phase=None,
+        normalize_phase=False,
+    ):
+        # Lightkurve v1.x assumed that `period` was given in days if no unit
+        # was specified. We maintain this behavior for backwards-compatibility.
+        if period is not None and not isinstance(period, Quantity):
+            period *= u.day
+        if epoch_time is not None and not isinstance(epoch_time, Time):
+            epoch_time = Time(
+                epoch_time, format=self.time.format, scale=self.time.scale
+            )
+        if (
+            epoch_phase is not None
+            and not isinstance(epoch_phase, Quantity)
+            and not normalize_phase
+        ):
+            epoch_phase *= u.day
+        if wrap_phase is not None and not isinstance(wrap_phase, Quantity):
+            wrap_phase *= u.day
+
+        # Warn if `epoch_time` appears to use the wrong format
+        if epoch_time is not None and epoch_time.value > 2450000:
+            if self.time.format == "bkjd":
+                warnings.warn(
+                    "`epoch_time` appears to be given in JD, "
+                    "however the light curve time uses BKJD "
+                    "(i.e. JD - 2454833).",
+                    LightkurveWarning,
+                )
+            elif self.time.format == "btjd":
+                warnings.warn(
+                    "`epoch_time` appears to be given in JD, "
+                    "however the light curve time uses BTJD "
+                    "(i.e. JD - 2457000).",
+                    LightkurveWarning,
+                )
+
+        ts = super().fold(
+            period=period,
+            epoch_time=epoch_time,
+            epoch_phase=epoch_phase,
+            wrap_phase=wrap_phase,
+            normalize_phase=normalize_phase,
+        )
+
+        # The folded time would pass the `TimeSeries` validation check if
+        # `normalize_phase=True`, so creating a `FoldedLightCurve` object
+        # requires the following three-step workaround:
+        # 1. Give the folded light curve a valid time column again
+        with ts._delay_required_column_checks():
+            folded_time = ts.time.copy()
+            ts.remove_column("time")
+            ts.add_column(self.time, name="time", index=0)
+        # 2. Create the folded object
+        lc = FoldedLightCurve(data=ts)
+        # 3. Restore the folded time
+        with lc._delay_required_column_checks():
+            lc.remove_column("time")
+            lc.add_column(folded_time, name="time", index=0)
+
+        # Add extra column and meta data specific to FoldedLightCurve
+        lc.add_column(
+            self.time.copy(), name="time_original", index=len(self._required_columns)
+        )
+        lc.meta["PERIOD"] = period
+        lc.meta["EPOCH_TIME"] = epoch_time
+        lc.meta["EPOCH_PHASE"] = epoch_phase
+        lc.meta["WRAP_PHASE"] = wrap_phase
+        lc.meta["NORMALIZE_PHASE"] = normalize_phase
+        lc.sort("time")
+
+        return lc
+
+    def normalize(self, unit="unscaled"):
+        validate_method(unit, ["unscaled", "percent", "ppt", "ppm"])
+        median_flux = np.nanmedian(self.flux)
+        std_flux = np.nanstd(self.flux)
+
+        # If the median flux is within half a standard deviation from zero, the
+        # light curve is likely zero-centered and normalization makes no sense.
+        if (median_flux == 0) or (
+            np.isfinite(std_flux) and (np.abs(median_flux) < 0.5 * std_flux)
+        ):
+            warnings.warn(
+                "The light curve appears to be zero-centered "
+                "(median={:.2e} +/- {:.2e}); `normalize()` will divide "
+                "the light curve by a value close to zero, which is "
+                "probably not what you want."
+                "".format(median_flux, std_flux),
+                LightkurveWarning,
+            )
+        # If the median flux is negative, normalization will invert the light
+        # curve and makes no sense.
+        if median_flux < 0:
+            warnings.warn(
+                "The light curve has a negative median flux ({:.2e});"
+                " `normalize()` will therefore divide by a negative "
+                "number and invert the light curve, which is probably"
+                "not what you want".format(median_flux),
+                LightkurveWarning,
+            )
+
+        # Create a new light curve instance and normalize its values
+        lc = self.copy()
+        lc.flux = lc.flux / median_flux
+        lc.flux_err = lc.flux_err / median_flux
+        if not lc.flux.unit:
+            lc.flux *= u.dimensionless_unscaled
+        if not lc.flux_err.unit:
+            lc.flux_err *= u.dimensionless_unscaled
+
+        # Set the desired relative (dimensionless) units
+        if unit == "percent":
+            lc.flux = lc.flux.to(u.percent)
+            lc.flux_err = lc.flux_err.to(u.percent)
+        elif unit in ("ppt", "ppm"):
+            lc.flux = lc.flux.to(unit)
+            lc.flux_err = lc.flux_err.to(unit)
+
+        lc.meta["NORMALIZED"] = True
+        return lc
+
+    def remove_nans(self, column: str = "flux"):
+        return self[~np.isnan(self[column])]  # This will return a sliced copy
+
+class FoldedLightCurve(LightCurve):
+
+    @property
+    def phase(self):
+        """Alias for `LightCurve.time`."""
+        return self.time
+
+    @property
+    def cycle(self):
+        cycle_epoch_start = self.epoch_time - self.period / 2
+        result = np.asarray(np.floor(((self.time_original - cycle_epoch_start) / self.period).value), dtype=int)
+        result = result - result.min()
+        return result
+
+    @property
+    def odd_mask(self):
+        return self.cycle % 2 == 1
+
+    @property
+    def even_mask(self):
+        return ~self.odd_mask
+
+    def _set_xlabel(self, kwargs):
+        if "xlabel" not in kwargs:
+            kwargs["xlabel"] = "Phase"
+            if isinstance(self.time, TimeDelta):
+                kwargs["xlabel"] += f" [{self.time.format.upper()}]"
+        return kwargs
+
+    def plot(self, **kwargs):
+        kwargs = self._set_xlabel(kwargs)
+        return super(FoldedLightCurve, self).plot(**kwargs)
+
+    def scatter(self, **kwargs):
+        kwargs = self._set_xlabel(kwargs)
+        return super(FoldedLightCurve, self).scatter(**kwargs)
+
+    def errorbar(self, **kwargs):
+        kwargs = self._set_xlabel(kwargs)
+        return super(FoldedLightCurve, self).errorbar(**kwargs)
+
+    def plot_river(self, **kwargs):
+        ax = super(FoldedLightCurve, self).plot_river(
+            period=self.period, epoch_time=self.epoch_time, **kwargs
+        )
+        return ax
+
 class KeplerLightCurve(LightCurve):
     """Subclass of :class:`LightCurve <lightkurve.lightcurve.LightCurve>`
     to represent data from NASA's Kepler and K2 mission."""
@@ -1418,3 +1619,5 @@ try:
 except registry.IORegistryError as exc:
     print(exc)
     pass  # necessary to enable autoreload during debugging
+
+