ExoIris 0.18.0__py3-none-any.whl → 0.19.2__py3-none-any.whl

exoiris/tsdata.py

@@ -15,13 +15,10 @@
 #  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 import warnings
-import numba
-
-import pandas as pd
 from collections.abc import Sequence
-
 from typing import Union, Optional
 
+import numba
 from astropy.io import fits as pf
 from astropy.stats import mad_std
 from astropy.utils import deprecated
@@ -29,15 +26,38 @@ from matplotlib.axes import Axes
 from matplotlib.figure import Figure
 from matplotlib.pyplot import subplots, setp
 from matplotlib.ticker import LinearLocator, FuncFormatter
-from numpy import isfinite, median, where, all, zeros_like, diff, asarray, interp, arange, floor, ndarray, \
-    ceil, newaxis, inf, array, ones, poly1d, polyfit, nanpercentile, atleast_2d, nan, linspace, any, sqrt, nanmedian
+from numpy import (
+    isfinite,
+    where,
+    all,
+    zeros_like,
+    diff,
+    asarray,
+    interp,
+    arange,
+    floor,
+    ndarray,
+    ceil,
+    newaxis,
+    inf,
+    array,
+    ones,
+    poly1d,
+    polyfit,
+    nanpercentile,
+    atleast_2d,
+    nan,
+    sqrt,
+    nanmedian,
+    nanmean,
+    unique,
+)
 from pytransit.orbits import fold
-from scipy.ndimage import median_filter
-from scipy.signal import medfilt
 
+from .binning import Binning, CompoundBinning
 from .ephemeris import Ephemeris
 from .util import bin2d
-from .binning import Binning, CompoundBinning
+
 
 class TSData:
     """
@@ -45,9 +65,10 @@ class TSData:
     fluxes, and errors. It provides methods for manipulating and analyzing the data.
     """
     def __init__(self, time: Sequence, wavelength: Sequence, fluxes: Sequence, errors: Sequence, name: str,
-                 noise_group: str = 'a', wl_edges : Sequence | None = None, tm_edges : Sequence | None = None,
+                 noise_group: int = 0, wl_edges : Sequence | None = None, tm_edges : Sequence | None = None,
                  transit_mask: ndarray | None = None, ephemeris: Ephemeris | None = None, n_baseline: int = 1,
-                 mask: ndarray = None, ephemeris_group: int = 0, offset_group: int = 0) -> None:
+                 mask: ndarray = None, epoch_group: int = 0, offset_group: int = 0,
+                 mask_nonfinite_errors: bool = True) -> None:
         """
         Parameters
         ----------
@@ -81,26 +102,42 @@ class TSData:
         if n_baseline < 1:
             raise ValueError("n_baseline must be greater than zero.")
 
-        if ephemeris_group < 0:
-            raise ValueError("ephemeris_group must be a non-negative integer.")
+        if noise_group < 0:
+            raise ValueError("noise_group must be a positive integer.")
+
+        if epoch_group < 0:
+            raise ValueError("epoch_group must be a non-negative integer.")
 
         if offset_group < 0:
             raise ValueError("offset_group must be a non-negative integer.")
 
+        if not all(isfinite(time)):
+            raise ValueError("The time array must contain only finite values.")
+
+        if not all(isfinite(wavelength)):
+            raise ValueError("The wavelength array must contain only finite values.")
+
         self.name: str = name
+        self.mask_nonfinite_errors: bool = mask_nonfinite_errors
         self.time: ndarray = time.copy()
         self.wavelength: ndarray = wavelength
-        self.mask: ndarray = mask if mask is not None else isfinite(fluxes) & isfinite(errors)
+        self.mask: ndarray = mask if mask is not None else isfinite(fluxes)
+        if self.mask_nonfinite_errors:
+            self.mask &= isfinite(errors)
         self.fluxes: ndarray = where(self.mask, fluxes, nan)
         self.errors: ndarray = where(self.mask, errors, nan)
         self.transit_mask: ndarray = transit_mask if transit_mask is not None else ones(time.size, dtype=bool)
-        self.ngid: int = 0
-        self.ephemeris: Ephemeris | None = ephemeris
+        self._ephemeris: Ephemeris | None = ephemeris
         self.n_baseline: int = n_baseline
-        self._noise_group: str = noise_group
-        self.ephemeris_group: int = ephemeris_group
+        self.noise_group: int = noise_group
+        self.epoch_group: int = epoch_group
         self.offset_group: int = offset_group
-        self._dataset: Optional['TSDataSet'] = None
+        self._dataset: Optional['TSDataGroup'] = None
+        self.minwl: float = 0.0
+        self.maxwl: float = inf
+        self.mintm: float = 0.0
+        self.maxtm: float = inf
+
         self._update()
 
         if wl_edges is None:
@@ -137,7 +174,7 @@ class TSData:
         mask = pf.ImageHDU(self.mask.astype(int), name=f'mask_{self.name}')
         data.header['ngroup'] = self.noise_group
         data.header['nbasel'] = self.n_baseline
-        data.header['epgroup'] = self.ephemeris_group
+        data.header['epgroup'] = self.epoch_group
         data.header['offgroup'] = self.offset_group
         #TODO: export ephemeris
         return pf.HDUList([time, wave, data, ootm, mask])
@@ -162,9 +199,21 @@ class TSData:
         data = hdul[f'DATA_{name}'].data.astype('d')
         ootm = hdul[f'OOTM_{name}'].data.astype(bool)
         mask = hdul[f'MASK_{name}'].data.astype(bool)
-        noise_group = hdul[f'DATA_{name}'].header['NGROUP']
-        ephemeris_group = hdul[f'DATA_{name}'].header['EPGROUP']
-        offset_group = hdul[f'DATA_{name}'].header['OFFGROUP']
+
+        try:
+            noise_group = hdul[f'DATA_{name}'].header['NGROUP']
+        except KeyError:
+            noise_group = 0
+
+        try:
+            ephemeris_group = hdul[f'DATA_{name}'].header['EPGROUP']
+        except KeyError:
+            ephemeris_group = 0
+
+        try:
+            offset_group = hdul[f'DATA_{name}'].header['OFFGROUP']
+        except KeyError:
+            offset_group = 0
 
         try:
             n_baseline = hdul[f'DATA_{name}'].header['NBASEL']
@@ -173,21 +222,29 @@ class TSData:
 
         #TODO: import ephemeris
         return TSData(time, wave, data[0], data[1], name=name, noise_group=noise_group, transit_mask=ootm,
-                      n_baseline=n_baseline, mask=mask, ephemeris_group=ephemeris_group, offset_group=offset_group)
+                      n_baseline=n_baseline, mask=mask, epoch_group=ephemeris_group, offset_group=offset_group)
 
     def __repr__(self) -> str:
         return f"TSData Name:'{self.name}' [{self.wavelength[0]:.2f} - {self.wavelength[-1]:.2f}] nwl={self.nwl} npt={self.npt}"
 
     @property
-    def noise_group(self) -> str:
-        """Noise group name."""
-        return self._noise_group
+    def ephemeris(self) -> Ephemeris:
+        """Ephemeris."""
+        return self._ephemeris
+
+    @ephemeris.setter
+    def ephemeris(self, ep: Ephemeris) -> None:
+        self._ephemeris = ep
+        self.mask_transit(ephemeris=ep)
 
-    @noise_group.setter
-    def noise_group(self, ng: str) -> None:
-        self._noise_group = ng
-        if self._dataset is not None:
-            self._dataset._update_nids()
+    @property
+    def bbox_wl(self) -> tuple[float, float]:
+        """Wavelength bounds of the bounding box."""
+        return self.minwl, self.maxwl
+
+    @property
+    def bbox_tm(self) -> tuple[float, float]:
+        return self.mintm, self.maxtm
 
     def mask_transit(self, t0: float | None = None, p: float | None = None, t14: float | None = None,
                      ephemeris : Ephemeris | None = None, elims: tuple[int, int] | None = None) -> 'TSData':
@@ -208,9 +265,9 @@ class TSData:
         """
         if (t0 and p and t14) or ephemeris is not None:
             if ephemeris is not None:
-                self.ephemeris = ephemeris
+                self._ephemeris = ephemeris
             else:
-                self.ephemeris = Ephemeris(t0, p, t14)
+                self._ephemeris = Ephemeris(t0, p, t14)
             phase = fold(self.time, self.ephemeris.period, self.ephemeris.zero_epoch)
             self.transit_mask = abs(phase) > 0.502 * self.ephemeris.duration
         elif elims is not None:
@@ -238,7 +295,19 @@ class TSData:
         """Update the internal attributes."""
         self.nwl = self.wavelength.size
         self.npt = self.time.size
-        self.wllims = self.wavelength.min(), self.wavelength.max()
+        self.minwl = self.wavelength.min()
+        self.maxwl = self.wavelength.max()
+        self.mintm = self.time.min()
+        self.maxtm = self.time.max()
+        if self._ephemeris is not None:
+            self.mask_transit(ephemeris=self._ephemeris)
+
+    def _update_data_mask(self) -> None:
+        self.mask = isfinite(self.fluxes)
+        if self.mask_nonfinite_errors:
+            self.mask &= isfinite(self.errors)
+        self.fluxes = where(self.mask, self.fluxes, nan)
+        self.errors = where(self.mask, self.errors, nan)
 
     def normalize_to_poly(self, deg: int = 1) -> 'TSData':
         """Normalize the baseline flux for each spectroscopic light curve.
@@ -266,11 +335,15 @@ class TSData:
                              "Call TSData.mask_transit(...) first.")
 
         for ipb in range(self.nwl):
-            bl = poly1d(polyfit(self.time[self.transit_mask & self.mask[ipb]],
-                                self.fluxes[ipb, self.transit_mask & self.mask[ipb]],
-                                deg=deg))(self.time)
-            self.fluxes[ipb, :] /= bl
-            self.errors[ipb, :] /= bl
+            mask = self.transit_mask & self.mask[ipb]
+            if mask.sum() > 2:
+                bl = poly1d(polyfit(self.time[mask], self.fluxes[ipb, mask], deg=deg))(self.time)
+                self.fluxes[ipb, :] /= bl
+                self.errors[ipb, :] /= bl
+            else:
+                self.fluxes[ipb, :] = nan
+                self.errors[ipb, :] = nan
+        self._update_data_mask()
         return self
 
     def normalize_to_median(self, s: slice) -> 'TSData':
@@ -286,7 +359,7 @@ class TSData:
         self.errors[:,:] /= n
         return self
 
-    def partition_time(self, tlims: tuple[tuple[float,float]]) -> 'TSDataSet':
+    def partition_time(self, tlims: tuple[tuple[float,float]]) -> 'TSDataGroup':
         """Partition the data into n segments defined by tlims.
 
         Parameters
@@ -299,20 +372,22 @@ class TSData:
         d = TSData(name=f'{self.name}_1', time=self.time[m], wavelength=self.wavelength,
                    fluxes=self.fluxes[:, m], errors=self.errors[:, m], mask=self.mask[:, m],
                    noise_group=self.noise_group,
-                   ephemeris_group=self.ephemeris_group,
+                   epoch_group=self.epoch_group,
                    offset_group=self.offset_group,
                    transit_mask=self.transit_mask[m],
                    ephemeris=self.ephemeris,
-                   n_baseline=self.n_baseline)
+                   n_baseline=self.n_baseline,
+                   mask_nonfinite_errors=self.mask_nonfinite_errors)
         for i, m in enumerate(masks[1:]):
             d = d + TSData(name=f'{self.name}_{i+2}', time=self.time[m], wavelength=self.wavelength,
                            fluxes=self.fluxes[:, m], errors=self.errors[:, m], mask=self.mask[:, m],
                            noise_group=self.noise_group,
-                           ephemeris_group=self.ephemeris_group,
+                           epoch_group=self.epoch_group,
                            offset_group=self.offset_group,
                            transit_mask=self.transit_mask[m],
                            ephemeris=self.ephemeris,
-                           n_baseline=self.n_baseline)
+                           n_baseline=self.n_baseline,
+                           mask_nonfinite_errors=self.mask_nonfinite_errors)
         return d
 
     def crop_wavelength(self, lmin: float, lmax: float, inplace: bool = True) -> 'TSData':
@@ -344,12 +419,13 @@ class TSData:
                           errors=self.errors[m],
                           mask=self.mask[m],
                           noise_group=self.noise_group,
-                          ephemeris_group=self.ephemeris_group,
+                          epoch_group=self.epoch_group,
                           offset_group=self.offset_group,
                           wl_edges=(self._wl_l_edges[m], self._wl_r_edges[m]),
                           tm_edges=(self._tm_l_edges, self._tm_r_edges),
                           transit_mask=self.transit_mask, ephemeris=self.ephemeris,
-                          n_baseline=self.n_baseline)
+                          n_baseline=self.n_baseline,
+                          mask_nonfinite_errors=self.mask_nonfinite_errors)
 
     def crop_time(self, tmin: float, tmax: float, inplace: bool = True) -> 'TSData':
         """Crop the data to include only the time range between lmin and lmax.
@@ -381,19 +457,20 @@ class TSData:
                           errors=self.errors[:, m],
                           mask = self.mask[:, m],
                           noise_group=self.noise_group,
-                          ephemeris_group=self.ephemeris_group,
+                          epoch_group=self.epoch_group,
                           offset_group=self.offset_group,
                           wl_edges=(self._wl_l_edges, self._wl_r_edges),
                           tm_edges=(self._tm_l_edges[m], self._tm_r_edges[m]),
                           transit_mask=self.transit_mask[m], ephemeris=self.ephemeris,
-                          n_baseline=self.n_baseline)
+                          n_baseline=self.n_baseline,
+                          mask_nonfinite_errors=self.mask_nonfinite_errors)
 
-    def remove_outliers(self, sigma: float = 5.0) -> 'TSData':
-        """Remove outliers along the wavelength axis.
+    # TODO: separate mask into bad data mask and outlier mask.
+    def mask_outliers(self, sigma: float = 5.0) -> 'TSData':
+        """Mask outliers along the wavelength axis.
 
-        Replace outliers along the wavelength axis with the value of a 5-point running median filter. Outliers are
-        defined as data points that deviate from the median by more than sigma times the median absolute deviation
-        along the wavelength axis.
+        Outliers are defined as data points that deviate from the running 5-point median by more
+        than sigma times the median absolute deviation along the wavelength axis.
 
         Parameters
         ----------
@@ -404,13 +481,18 @@ class TSData:
         ----
         The data will be modified in place.
         """
-        fm = median(self.fluxes, axis=0)
-        fe = mad_std(self.fluxes, axis=0)
+        fm = nanmedian(self.fluxes, axis=0)
+        fe = mad_std(self.fluxes, axis=0, ignore_nan=True)
         self.mask &= abs(self.fluxes - fm) / fe < sigma
         self.fluxes = where(self.mask, self.fluxes, nan)
         self.errors = where(self.mask, self.errors, nan)
         return self
 
+    @deprecated("0.10", alternative="TSData.mask_outliers")
+    def remove_outliers(self, sigma: float = 5.0) -> 'TSData':
+        """Remove outliers along the wavelength axis."""
+        self.mask_outliers(sigma=sigma)
+
     def plot(self, ax=None, vmin: float = None, vmax: float = None, cmap=None, figsize=None, data=None,
              plims: tuple[float, float] | None = None) -> Figure:
         """Plot the spectroscopic light curves as a 2D image.
@@ -510,7 +592,7 @@ class TSData:
             fig = ax.figure
         tref = floor(self.time.min())
 
-        ax.plot(self.time, self.fluxes.mean(0))
+        ax.plot(self.time, nanmean(self.fluxes, 0))
         if self.ephemeris is not None:
             [ax.axvline(tl, ls='--', c='k') for tl in self.ephemeris.transit_limits(self.time.mean())]
 
@@ -547,7 +629,7 @@ class TSData:
         """
         return self.plot(ax=ax, figsize=figsize, data=where(self.transit_mask, self.fluxes, nan))
 
-    def __add__(self, other: Union['TSData', 'TSDataSet']) -> 'TSDataSet':
+    def __add__(self, other: Union['TSData', 'TSDataGroup']) -> 'TSDataGroup':
         """Combine two transmission spectra along the wavelength axis.
 
         Parameters
@@ -557,12 +639,12 @@ class TSData:
 
         Returns
         -------
-        TSDataSet
+        TSDataGroup
         """
         if isinstance(other, TSData):
-            return TSDataSet([self, other])
+            return TSDataGroup([self, other])
         else:
-            return TSDataSet([self]) + other
+            return TSDataGroup([self]) + other
 
     def bin_wavelength(self, binning: Optional[Union[Binning, CompoundBinning]] = None,
                        nb: Optional[int] = None, bw: Optional[float] = None, r: Optional[float] = None,
@@ -592,7 +674,7 @@ class TSData:
         with warnings.catch_warnings():
             warnings.simplefilter('ignore', numba.NumbaPerformanceWarning)
             if binning is None:
-                binning = Binning(self.wllims[0], self.wllims[1], nb=nb, bw=bw, r=r)
+                binning = Binning(self.bbox_wl[0], self.bbox_wl[1], nb=nb, bw=bw, r=r)
             bf, be = bin2d(self.fluxes, self.errors, self._wl_l_edges, self._wl_r_edges,
                            binning.bins, estimate_errors=estimate_errors)
             if not all(isfinite(be)):
@@ -602,7 +684,7 @@ class TSData:
                           name=self.name,
                           tm_edges=(self._tm_l_edges, self._tm_r_edges),
                           noise_group=self.noise_group,
-                          ephemeris_group=self.ephemeris_group,
+                          epoch_group=self.epoch_group,
                           offset_group=self.offset_group,
                           transit_mask=self.transit_mask,
                           ephemeris=self.ephemeris,
@@ -644,14 +726,14 @@ class TSData:
                        noise_group=self.noise_group,
                        ephemeris=self.ephemeris,
                        n_baseline=self.n_baseline,
-                       ephemeris_group=self.ephemeris_group,
+                       epoch_group=self.epoch_group,
                        offset_group=self.offset_group)
             if self.ephemeris is not None:
                 d.mask_transit(ephemeris=self.ephemeris)
             return d
 
-class TSDataSet:
-    """`TSDataSet` is a high-level data storage class that can contain multiple `TSData` objects.
+class TSDataGroup:
+    """`TSDataGroup` is a high-level data storage class that can contain multiple `TSData` objects.
     """
     def __init__(self, data: Sequence[TSData]):
         self.data: list[TSData] = []
@@ -659,7 +741,7 @@ class TSDataSet:
         self.wlmax: float = -inf
         self.tmin: float = inf
         self.tmax: float = -inf
-        self.ngids: ndarray = array([])
+        self._noise_groups: array | None = None
         for d in data:
             self._add_data(d)
 
@@ -668,19 +750,12 @@ class TSDataSet:
             raise ValueError('A TSData object with the same name already exists.')
         d._dataset = self
         self.data.append(d)
-        self._update_nids()
+        self._noise_groups = array([d.noise_group for d in self.data])
         self.wlmin = min(self.wlmin, d.wavelength.min())
         self.wlmax = max(self.wlmax, d.wavelength.max())
         self.tmin = min(self.tmin, d.time.min())
         self.tmax = max(self.tmax, d.time.max())
 
-    def _update_nids(self):
-        ngs =  pd.Categorical(self.noise_groups)
-        self.unique_noise_groups = list(ngs.categories)
-        self.ngids = ngs.codes.astype(int)
-        for i,d in enumerate(self.data):
-            d.ngid = self.ngids[i]
-
     @property
     def names(self) -> list[str]:
         """List of data set names."""
@@ -707,20 +782,25 @@ class TSDataSet:
         return [d.errors for d in self.data]
 
     @property
-    def noise_groups(self) -> list[str]:
-        """List of noise group names."""
-        return [d.noise_group for d in self.data]
+    def noise_groups(self) -> ndarray[int]:
+        """Array of noise groups."""
+        return self._noise_groups
 
     @property
     def n_noise_groups(self) -> int:
         """Number of noise groups."""
-        return len(set(self.noise_groups))
+        return len(unique(self.noise_groups))
 
     @property
     def offset_groups(self) -> list[int]:
         """List of offset groups."""
         return [d.offset_group for d in self.data]
 
+    @property
+    def epoch_groups(self) -> list[int]:
+        """List of epoch groups."""
+        return [d.epoch_group for d in self.data]
+
     @property
     def n_baselines(self) -> list[int]:
         """Number of baseline coefficients for each data set."""
@@ -749,7 +829,7 @@ class TSDataSet:
         return hdul
 
     @staticmethod
-    def import_fits(hdul: pf.HDUList) -> 'TSDataSet':
+    def import_fits(hdul: pf.HDUList) -> 'TSDataGroup':
         """Import all the data from a FITS HDU list.
 
         Parameters
@@ -759,14 +839,14 @@ class TSDataSet:
 
         Returns
         -------
-        TSDataSet
+        TSDataGroup
         """
         ds = hdul['DATASET']
         data = []
         for i in range(ds.header['NDATA']):
             name = ds.header[f'NAME_{i}']
             data.append(TSData.import_fits(name, hdul))
-        return TSDataSet(data)
+        return TSDataGroup(data)
 
     def mask_transit(self, tc: float, p: float, t14: float):
         for d in self.data:
@@ -779,7 +859,7 @@ class TSDataSet:
         return self.size
 
     def __repr__(self):
-        return f"TSDataSet with {self.size} groups"
+        return f"TSDataGroup with {self.size} groups"
 
     def plot(self, axs=None, vmin: float = None, vmax: float = None, ncols: int = 1, cmap=None, figsize=None, data: ndarray | None = None) -> Figure:
         """Plot all the data sets.
@@ -849,6 +929,10 @@ class TSDataSet:
 
     def __add__(self, other):
         if isinstance(other, TSData):
-            return TSDataSet(self.data + [other])
-        elif isinstance(other, TSDataSet):
-            return TSDataSet(self.data + other.data)
+            return TSDataGroup(self.data + [other])
+        elif isinstance(other, TSDataGroup):
+            return TSDataGroup(self.data + other.data)
+
+
+class TSDataSet(TSDataGroup):
+    pass