- Introduced epoch_group to replace ephemeris_group for better clarity.

- Added `mask_nonfinite_errors` flag to handle masking invalid errors in data.
- Enhanced transit parameters by separating orbital and epoch-related data.
- Updated outlier removal logic to a dedicated `mask_outliers` method.
- Improved initialization and handling of transit center parameters.
- Deprecated `remove_outliers` in favor of `mask_outliers`.
- Refactored `nanmean` and other NumPy operations for improved efficiency.

Author: hpparvi

exoiris/tsdata.py

@@ -29,8 +29,9 @@
 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 any, 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, \
+    nanmean
 from pytransit.orbits import fold
 from scipy.ndimage import median_filter
 from scipy.signal import medfilt
@@ -45,9 +46,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,8 +83,11 @@ def __init__(self, time: Sequence, wavelength: Sequence, fluxes: Sequence, error
         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.")
@@ -94,17 +99,20 @@ def __init__(self, time: Sequence, wavelength: Sequence, fluxes: Sequence, error
             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.epoch_group: int = epoch_group
         self.offset_group: int = offset_group
         self._dataset: Optional['TSDataSet'] = None
         self._update()
@@ -143,7 +151,7 @@ def export_fits(self) -> pf.HDUList:
         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])
@@ -191,7 +199,7 @@ def import_fits(name: str, hdul: pf.HDUList) -> '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}"
@@ -207,6 +215,16 @@ def noise_group(self, ng: str) -> None:
         if self._dataset is not None:
             self._dataset._update_nids()
 
+    @property
+    def ephemeris(self) -> Ephemeris:
+        """Ephemeris."""
+        return self._ephemeris
+
+    @ephemeris.setter
+    def ephemeris(self, ep: Ephemeris) -> None:
+        self._ephemeris = ep
+        self.mask_transit(ephemeris=ep)
+
     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':
         """Create a transit mask based on a given ephemeris or exposure index limits.
@@ -226,9 +244,9 @@ def mask_transit(self, t0: float | None = None, p: float | None = None, t14: flo
         """
         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:
@@ -257,6 +275,15 @@ def _update(self) -> None:
         self.nwl = self.wavelength.size
         self.npt = self.time.size
         self.wllims = self.wavelength.min(), self.wavelength.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.
@@ -289,6 +316,7 @@ def normalize_to_poly(self, deg: int = 1) -> 'TSData':
                                 deg=deg))(self.time)
             self.fluxes[ipb, :] /= bl
             self.errors[ipb, :] /= bl
+        self._update_data_mask()
         return self
 
     def normalize_to_median(self, s: slice) -> 'TSData':
@@ -317,20 +345,22 @@ def partition_time(self, tlims: tuple[tuple[float,float]]) -> 'TSDataSet':
         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':
@@ -362,12 +392,13 @@ def crop_wavelength(self, lmin: float, lmax: float, inplace: bool = True) -> 'TS
                           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.
@@ -399,19 +430,20 @@ def crop_time(self, tmin: float, tmax: float, inplace: bool = True) -> '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
         ----------
@@ -422,13 +454,18 @@ def remove_outliers(self, sigma: float = 5.0) -> '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.
@@ -528,7 +565,7 @@ def plot_white(self, ax: Axes | None = None, figsize: tuple[float, float] | None
             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())]
 
@@ -620,7 +657,7 @@ def bin_wavelength(self, binning: Optional[Union[Binning, CompoundBinning]] = No
                           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,
@@ -662,7 +699,7 @@ def bin_time(self, binning: Optional[Union[Binning, CompoundBinning]] = None,
                        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)
@@ -739,6 +776,11 @@ 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."""

exoiris/tslpf.py

@@ -203,6 +203,7 @@ def _init_parameters(self) -> None:
         self.ps = ParameterSet([])
         self._init_p_star()
         self._init_p_orbit()
+        self._init_p_transit_centers()
         self._init_p_limb_darkening()
         self._init_p_radius_ratios()
         self._init_p_noise()
@@ -311,15 +312,22 @@ def _init_p_limb_darkening(self) -> None:
 
     def _init_p_orbit(self):
         ps = self.ps
-        pp = [GParameter('tc', 'zero_epoch', '', NP(0.0, 0.1), (-inf, inf)),
-              GParameter('p', 'period', 'd', NP(1.0, 1e-5), (0, inf)),
+        pp = [GParameter('p', 'period', 'd', NP(1.0, 1e-5), (0, inf)),
               GParameter('b', 'impact_parameter', 'R_s', UP(0.0, 1.0), (0, inf)),
               GParameter('secw', 'sqrt(e) cos(w)', '', NP(0.0, 1e-5), (-1, 1)),
               GParameter('sesw', 'sqrt(e) sin(w)', '', NP(0.0, 1e-5), (-1, 1))]
         ps.add_global_block('orbit', pp)
         self._start_orbit = ps.blocks[-1].start
         self._sl_orbit = ps.blocks[-1].slice
 
+    def _init_p_transit_centers(self):
+        ps = self.ps
+        neps = max(self.data.epoch_groups) + 1
+        pp = [GParameter(f'tc_{i:02d}', f'zero_epoch_{i:02d}', '', NP(0.0, 0.1), (-inf, inf)) for i in range(neps)]
+        ps.add_global_block('transit_centers', pp)
+        self._start_tcs = ps.blocks[-1].start
+        self._sl_tcs = ps.blocks[-1].slice
+
     def _init_p_radius_ratios(self):
         ps = self.ps
         pp = [GParameter(f'k_{k:08.5f}', fr'radius ratio at {k:08.5f} $\mu$m', 'A_s', UP(0.02, 0.2), (0, inf)) for k in self.k_knots]
@@ -592,22 +600,24 @@ def transit_model(self, pv, copy=True):
         """
         pv = atleast_2d(pv)
         ldp = self._eval_ldc(pv)
-        t0 = pv[:, 1]
-        p = pv[:, 2]
+        t0s = pv[:, self._sl_tcs]
         k = self._eval_k(pv[:, self._sl_rratios])
+        p = pv[:, 1]
         aor = as_from_rhop(pv[:, 0], p)
-        inc = i_from_ba(pv[:, 3], aor)
-        ecc = pv[:, 4] ** 2 + pv[:, 5] ** 2
-        w = arctan2(pv[:, 5], pv[:, 4])
+        inc = i_from_ba(pv[:, 2], aor)
+        ecc = pv[:, 3] ** 2 + pv[:, 4] ** 2
+        w = arctan2(pv[:, 4], pv[:, 3])
+        epids = self.data.epoch_groups
         fluxes = []
         if isinstance(self.ldmodel, LDTkLD):
             ldp, istar = self.ldmodel(self.tms[0].mu, ldp)
             ldpi = dstack([ldp, istar])
             for i, tm in enumerate(self.tms):
-                fluxes.append(tm.evaluate(k[i], ldpi[:, self.ldmodel.wlslices[i], :], t0, p, aor, inc, ecc, w, copy))
+                fluxes.append(tm.evaluate(k[i], ldpi[:, self.ldmodel.wlslices[i], :],
+                                          t0s[:, epids[i]], p, aor, inc, ecc, w, copy))
         else:
             for i, tm in enumerate(self.tms):
-                fluxes.append(tm.evaluate(k[i], ldp[i], t0, p, aor, inc, ecc, w, copy))
+                fluxes.append(tm.evaluate(k[i], ldp[i], t0s[:, epids[i]], p, aor, inc, ecc, w, copy))
 
         for i, d in enumerate(self.data):
             if d.offset_group > 0: