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

exoiris/__init__.py

@@ -15,7 +15,7 @@
 #  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 from .exoiris import ExoIris, load_model # noqa
-from .tsdata import TSData, TSDataSet # noqa
+from .tsdata import TSData, TSDataGroup # noqa
 from .binning import Binning # noqa
 from .ldtkld import LDTkLD # noqa
 from .tslpf import clean_knots # noqa

exoiris/exoiris.py

@@ -24,27 +24,24 @@ from typing import Optional, Callable, Any, Literal
 
 import astropy.io.fits as pf
 import astropy.units as u
-import emcee
+import matplotlib.axes
 import pandas as pd
-import pytransit.utils.de
-import seaborn as sb
 from astropy.table import Table
 from celerite2 import GaussianProcess, terms
 from emcee import EnsembleSampler
 from matplotlib.pyplot import subplots, setp, figure, Figure, Axes
-from numpy import (where, sqrt, clip, percentile, median, squeeze, floor, ndarray,
-                   array, inf, newaxis, arange, tile, sort, argsort, concatenate, full, nan, r_, nanpercentile, log10,
-                   ceil)
-from numpy.random import normal, permutation
+from numpy import (any, where, sqrt, clip, percentile, median, squeeze, floor, ndarray, isfinite,
+                   array, inf, arange, argsort, concatenate, full, nan, r_, nanpercentile, log10,
+                   ceil, unique)
+from numpy.random import normal
 from pytransit import UniformPrior, NormalPrior
-from pytransit.orbits import epoch
 from pytransit.param import ParameterSet
 from pytransit.utils.de import DiffEvol
 from scipy.stats import norm
 from uncertainties import UFloat
 
 from .ldtkld import LDTkLD
-from .tsdata import TSData, TSDataSet
+from .tsdata import TSData, TSDataGroup
 from .tslpf import TSLPF
 from .wlpf import WhiteLPF
 
@@ -73,7 +70,7 @@ def load_model(fname: Path | str, name: str | None = None):
         If the file format is invalid or does not match the expected format.
     """
     with pf.open(fname) as hdul:
-        data = TSDataSet.import_fits(hdul)
+        data = TSDataGroup.import_fits(hdul)
 
         if hdul[0].header['LDMODEL'] == 'ldtk':
             filters, teff, logg, metal, dataset = pickle.loads(codecs.decode(json.loads(hdul[0].header['LDTKLD']).encode(), "base64"))
@@ -86,11 +83,28 @@ def load_model(fname: Path | str, name: str | None = None):
         except KeyError:
             ip = 'bspline'
 
-        #TODO: save and load the noise model information
-        a = ExoIris(name or hdul[0].header['NAME'], ldmodel=ldm, data=data, interpolation=ip)
+        try:
+            noise_model = hdul[0].header['NOISE']
+        except KeyError:
+            noise_model = "white"
+
+        a = ExoIris(name or hdul[0].header['NAME'], ldmodel=ldm, data=data, noise_model=noise_model, interpolation=ip)
         a.set_radius_ratio_knots(hdul['K_KNOTS'].data.astype('d'))
         a.set_limb_darkening_knots(hdul['LD_KNOTS'].data.astype('d'))
 
+        # Read the white light curve models if they exist.
+        try:
+            tb = Table.read(hdul['WHITE_DATA'])
+            white_ids = tb['id'].data
+            uids = unique(white_ids)
+            a._white_times = [tb['time'].data[white_ids == i] for i in uids]
+            a._white_fluxes = [tb['flux_obs'].data[white_ids == i] for i in uids]
+            a._white_errors = [tb['flux_obs_err'].data[white_ids == i] for i in uids]
+            a._white_models = [tb['flux_mod'].data[white_ids == i] for i in uids]
+
+        except KeyError:
+            pass
+
         try:
             a.period = hdul[0].header['P']
             a.zero_epoch = hdul[0].header['T0']
@@ -116,8 +130,8 @@ class ExoIris:
     """The core ExoIris class providing tools for exoplanet transit spectroscopy.
     """
 
-    def __init__(self, name: str, ldmodel, data: TSDataSet | TSData, nk: int = 50, nldc: int = 10, nthreads: int = 1,
-                 tmpars: dict | None = None, noise_model: str = 'white',
+    def __init__(self, name: str, ldmodel, data: TSDataGroup | TSData, nk: int = 50, nldc: int = 10, nthreads: int = 1,
+                 tmpars: dict | None = None, noise_model: Literal["white", "fixed_gp", "free_gp"] = 'white',
                  interpolation: Literal['bspline', 'pchip', 'makima'] = 'bspline'):
         """
         Parameters
@@ -139,10 +153,31 @@ class ExoIris:
         noise_model
             The noise model to use. Should be either "white" for white noise or "fixed_gp" for Gaussian Process.
         """
-        data = TSDataSet([data]) if isinstance(data, TSData) else data
-        self._tsa: TSLPF = TSLPF(name, ldmodel, data, nk=nk, nldc=nldc, nthreads=nthreads, tmpars=tmpars,
+        data = TSDataGroup([data]) if isinstance(data, TSData) else data
+
+        for d in data:
+            if any(~isfinite(d.fluxes[d.mask])):
+                raise ValueError(f"The {d.name} data set flux array contains unmasked noninfinite values.")
+
+            if any(~isfinite(d.errors[d.mask])):
+                raise ValueError(f"The {d.name} data set error array contains unmasked noninfinite values.")
+
+        ngs = array(data.noise_groups)
+        if not ((ngs.min() == 0) and (ngs.max() + 1 == unique(ngs).size)):
+            raise ValueError("The noise groups must start from 0 and be consecutive.")
+
+        ogs = array(data.offset_groups)
+        if not ((ogs.min() == 0) and (ogs.max() + 1 == unique(ogs).size)):
+            raise ValueError("The offset groups must start from 0 and be consecutive.")
+
+        egs = array(data.epoch_groups)
+        if not ((egs.min() == 0) and (egs.max() + 1 == unique(egs).size)):
+            raise ValueError("The epoch groups must start from 0 and be consecutive.")
+
+        self._tsa = TSLPF(self, name, ldmodel, data, nk=nk, nldc=nldc, nthreads=nthreads, tmpars=tmpars,
                                  noise_model=noise_model, interpolation=interpolation)
-        self._wa: WhiteLPF | None = None
+        self._wa: None | WhiteLPF = None
+
         self.nthreads: int = nthreads
 
         self.period: float | None = None
@@ -150,6 +185,11 @@ class ExoIris:
         self.transit_duration: float | None= None
         self._tref = floor(self.data.tmin)
 
+        self._white_times: None | list[ndarray] = None
+        self._white_fluxes: None | list[ndarray] = None
+        self._white_errors: None | list[ndarray] = None
+        self._white_models: None | list[ndarray] = None
+
     def lnposterior(self, pvp: ndarray) -> ndarray:
         """Calculate the log posterior probability for a single parameter vector or an array of parameter vectors.
 
@@ -180,7 +220,7 @@ class ExoIris:
         """
         self._tsa.set_noise_model(noise_model)
 
-    def set_data(self, data: TSData | TSDataSet) -> None:
+    def set_data(self, data: TSData | TSDataGroup) -> None:
         """Set the model data.
 
         Parameters
@@ -188,7 +228,7 @@ class ExoIris:
         data
            The spectroscopic transit light curve.
         """
-        data = TSDataSet([data]) if isinstance(data, TSData) else data
+        data = TSDataGroup([data]) if isinstance(data, TSData) else data
         self._tsa.set_data(data)
 
     def set_prior(self, parameter: Literal['radius ratios', 'baselines', 'wn multipliers'] | str,
@@ -279,7 +319,7 @@ class ExoIris:
         metal = (metal.n, metal.s) if isinstance(metal, UFloat) else metal
         self._tsa.set_ldtk_prior(teff, logg, metal, dataset, width, uncertainty_multiplier)
 
-    def set_gp_hyperparameters(self, sigma: float, rho: float) -> None:
+    def set_gp_hyperparameters(self, sigma: float, rho: float, idata: None | int = None) -> None:
         """Set Gaussian Process (GP) hyperparameters assuming a Matern-3/2 kernel.
 
         Parameters
@@ -288,8 +328,10 @@ class ExoIris:
             The kernel amplitude parameter.
         rho
             The length scale parameter.
+        idata
+            The data set for which to set the hyperparameters. If None, the hyperparameters are set for all data sets.
         """
-        self._tsa.set_gp_hyperparameters(sigma, rho)
+        self._tsa.set_gp_hyperparameters(sigma, rho, idata)
 
     def set_gp_kernel(self, kernel: terms.Term) -> None:
         """Set the Gaussian Process (GP) kernel.
@@ -311,7 +353,7 @@ class ExoIris:
         self._tsa.name = name
 
     @property
-    def data(self) -> TSDataSet:
+    def data(self) -> TSDataGroup:
         """Analysis data set."""
         return self._tsa.data
 
@@ -378,22 +420,35 @@ class ExoIris:
     @property
     def white_times(self) -> list[ndarray]:
         """White light curve time arrays."""
-        return self._wa.times
+        if self._wa is None:
+            return self._white_times
+        else:
+            return self._wa.times
 
     @property
     def white_fluxes(self) -> list[ndarray]:
         """White light curve flux arrays."""
-        return self._wa.fluxes
+        if self._wa is None:
+            return self._white_fluxes
+        else:
+            return self._wa.fluxes
 
     @property
     def white_models(self) -> list[ndarray]:
-        fm = self._wa.flux_model(self._wa._local_minimization.x)
-        return [fm[sl] for sl in self._wa.lcslices]
+        """Fitted white light curve flux model arrays."""
+        if self._wa is None:
+            return self._white_models
+        else:
+            fm = self._wa.flux_model(self._wa._local_minimization.x)
+            return [fm[sl] for sl in self._wa.lcslices]
 
     @property
     def white_errors(self) -> list[ndarray]:
         """White light curve flux error arrays."""
-        return self._wa.std_errors
+        if self._wa is None:
+            return self._white_errors
+        else:
+            return self._wa.std_errors
 
     def add_radius_ratio_knots(self, knot_wavelengths: Sequence) -> None:
         """Add radius ratio (k) knots.
@@ -464,65 +519,36 @@ class ExoIris:
         """Print the model parameterization."""
         self._tsa.print_parameters(1)
 
-    def plot_setup(self, figsize: tuple[float, float] | None =None, xscale: str | None = None, xticks: Sequence | None = None) -> Figure:
+    def plot_setup(self, figsize: tuple[float, float] | None = None,
+                   ax: matplotlib.axes.Axes | None = None,
+                   xscale: str | None = None, xticks: Sequence | None = None,
+                   yshift: float = 0.1, mh:float = 0.08, side_margin: float = 0.05,
+                   lw: float = 0.5, c='k') -> Figure:
         """Plot the model setup with limb darkening knots, radius ratio knots, and data binning.
-
-        Parameters
-        ----------
-        figsize
-            The size of the figure in inches.
-        xscale
-            The scale of the x-axis. If provided, the x-axis scale of all the subplots will be set to this value.
-        xticks
-            The list of x-axis tick values for all the subplots. If provided, the x-axis ticks of all the subplots will
-            be set to these values.
-
-        Returns
-        -------
-        Figure
-            The matplotlib Figure object that contains the created subplots.
-
         """
-        using_ldtk = isinstance(self._tsa.ldmodel, LDTkLD)
+        if ax is None:
+            fig, ax = subplots(figsize=figsize, constrained_layout=True)
+        else:
+            fig = ax.figure
 
-        if not using_ldtk:
-            figsize = figsize or (13, 4)
-            fig, axs = subplots(3, 1, figsize=figsize, sharex='all', sharey='all')
-            axl, axk, axw = axs
+        ndata = self.data.size
 
-            axl.vlines(self._tsa.ld_knots, 0.1, 0.5, ec='k')
-            axl.text(0.01, 0.90, 'Limb darkening knots', va='top', transform=axl.transAxes)
-        else:
-            figsize = figsize or (13, 2*4/3)
-            fig, axs = subplots(2, 1, figsize=figsize, sharex='all', sharey='all')
-            axk, axw = axs
-            axl = None
-
-        axk.vlines(self._tsa.k_knots, 0.1, 0.5, ec='k')
-        axk.text(0.01, 0.90, 'Radius ratio knots', va='top', transform=axk.transAxes)
-        for ds in self.data:
-            axw.vlines(ds.wavelength, 0.1, 0.5, ec='k')
-        axw.text(0.01, 0.90, 'Wavelength bins', va='top', transform=axw.transAxes)
-
-        if not using_ldtk:
-            sb.despine(ax=axl, top=False, bottom=True, right=False)
-            sb.despine(ax=axk, top=True, bottom=True, right=False)
-        else:
-            sb.despine(ax=axk, top=False, bottom=True, right=False)
+        for i, d in enumerate(self.data):
+            ax.vlines(d.wavelength, ymin=i*yshift, ymax=i*yshift+mh, colors=c, lw=lw)
 
-        sb.despine(ax=axw, top=True, bottom=False, right=False)
-        setp(axs, xlim=(self.data.wlmin-0.02, self.data.wlmax+0.02), yticks=[], ylim=(0, 0.9))
-        setp(axw, xlabel=r'Wavelength [$\mu$m]')
-        setp(axs[0].get_xticklines(), visible=False)
-        setp(axs[0].get_xticklabels(), visible=False)
-        setp(axs[1].get_xticklines(), visible=False)
-        setp(axs[-1].get_xticklines(), visible=True)
+        i = ndata + 1
+        ax.vlines(self._tsa.ld_knots, ymin=i*yshift, ymax=i*yshift+mh, colors=c, lw=lw)
+
+        i = ndata + 3
+        ax.vlines(self.k_knots, ymin=i*yshift, ymax=i*yshift+mh, colors=c, lw=lw)
 
         if xscale:
-            setp(axs, xscale=xscale)
+            setp(ax, xscale=xscale)
         if xticks is not None:
-            [ax.set_xticks(xticks, labels=xticks) for ax in axs]
-        fig.tight_layout()
+            ax.set_xticks(xticks, labels=xticks)
+
+        setp(ax, yticks=[], xlim=(self.data.wlmin-side_margin, self.data.wlmax+side_margin), xlabel=r'Wavelength [$\mu$m]')
+        ax.set_yticks(concatenate([arange(ndata), arange(ndata+1, ndata+4, 2)])*yshift+0.5*mh, labels=[n.replace("_", " ") for n in self.data.names] + ["Limb darkening knots", "Radius ratio knots"])
         return fig
 
     def fit_white(self, niter: int = 500) -> None:
@@ -537,10 +563,10 @@ class ExoIris:
         self._wa.optimize_global(niter, plot_convergence=False, use_tqdm=False)
         self._wa.optimize()
         pv = self._wa._local_minimization.x
-        self.period = pv[1]
+        self.period = pv[0]
         self.zero_epoch = self._wa.transit_center
         self.transit_duration = self._wa.transit_duration
-        self.data.mask_transit(pv[0], pv[1], self.transit_duration)
+        self.data.mask_transit(self.zero_epoch, self.period, self.transit_duration)
 
     def plot_white(self, axs=None, figsize: tuple[float, float] | None = None, ncols: int | None=None) -> Figure:
         """Plot the white light curve data with the best-fit model.
@@ -556,7 +582,25 @@ class ExoIris:
         """
         return self._wa.plot(axs=axs, figsize=figsize, ncols=ncols or min(self.data.size, 2))
 
-    def plot_white_gp_predictions(self, axs = None, ncol: int = 1, figsize: tuple[float, float] | None = None):
+    def plot_white_gp_predictions(self, axs = None, ncol: int = 1, figsize: tuple[float, float] | None = None) -> None:
+        """Plot the predictions of a Gaussian Process model for white light curves and residuals.
+
+        Parameters
+        ----------
+        axs
+            Axes array in which the plots are drawn. If None, new subplots are generated.
+        ncol
+            The number of columns for the created grid of subplots. Used only if axs is None.
+        figsize
+            Size of the figure in inches (width, height). Used only if axs is None. If None,
+            a default size is used.
+
+        Notes
+        -----
+        The number of rows for the subplots is determined dynamically based on the shape of
+        the data and the number of columns specified (ncol). If the provided axes array (axs)
+        does not accommodate all the subplots, the behavior is undefined.
+        """
         ndata = self.data.size
 
         if axs is None:
@@ -657,12 +701,18 @@ class ExoIris:
 
                 pv0 = self._wa._local_minimization.x
                 x0 = self._tsa.ps.sample_from_prior(npop)
-                x0[:, 0] = normal(pv0[2], 0.05, size=npop)
-                x0[:, 1] = normal(pv0[0], 1e-4, size=npop)
-                x0[:, 2] = normal(pv0[1], 1e-5, size=npop)
-                x0[:, 3] = clip(normal(pv0[3], 0.01, size=npop), 0.0, 1.0)
+                x0[:, 0] = clip(normal(pv0[1], 0.05, size=npop), 0.01, inf)
+                x0[:, 1] = clip(normal(pv0[0], 1e-4, size=npop), 0.01, inf)
+                x0[:, 2] = clip(normal(pv0[2], 1e-3, size=npop), 0.0, 1.0)
+
+                nep = max(self.data.epoch_groups) + 1
+                for i in range(nep):
+                    pida = self.ps.find_pid(f'tc_{i:02d}')
+                    pidb = self._wa.ps.find_pid(f'tc_{i:02d}')
+                    x0[:, pida] = normal(pv0[pidb], 0.001, size=npop)
+
                 sl = self._tsa._sl_rratios
-                x0[:, sl] = normal(sqrt(pv0[4]), 0.001, size=(npop, self.nk))
+                x0[:, sl] = normal(sqrt(pv0[self._wa.ps.find_pid('k2')]), 0.001, size=(npop, self.nk))
                 for i in range(sl.start, sl.stop):
                     x0[:, i] = clip(x0[:, i], 1.001*self.ps[i].prior.a, 0.999*self.ps[i].prior.b)
 
@@ -1043,6 +1093,7 @@ class ExoIris:
         pri.header['t14'] = self.transit_duration
         pri.header['ndgroups'] = self.data.size
         pri.header['interp'] = self._tsa.interpolation
+        pri.header['noise'] = self._tsa.noise_model
 
         pr = pf.ImageHDU(name='priors')
         priors = [pickle.dumps(p) for p in self.ps]
@@ -1061,6 +1112,34 @@ class ExoIris:
         hdul = pf.HDUList([pri, k_knots, ld_knots, pr])
         hdul += self.data.export_fits()
 
+        if self._wa is not None and self._wa._local_minimization is not None:
+            wa_data = pf.BinTableHDU(
+                Table(
+                    [
+                        self._wa.lcids,
+                        self._wa.timea,
+                        concatenate(self.white_models),
+                        self._wa.ofluxa,
+                        concatenate(self._wa.std_errors),
+                    ],
+                    names="id time flux_mod flux_obs flux_obs_err".split(),
+                ), name='white_data'
+            )
+            hdul.append(wa_data)
+
+            names = []
+            counts = {}
+            for p in self._wa.ps.names:
+                if p not in counts.keys():
+                    counts[p] = 0
+                    names.append(p)
+                else:
+                    counts[p] += 1
+                    names.append(f'{p}_{counts[p]}')
+
+            wa_params = pf.BinTableHDU(Table(self._wa._local_minimization.x, names=names), name='white_params')
+            hdul.append(wa_params)
+
         if self._tsa.de is not None:
             de = pf.BinTableHDU(Table(self._tsa._de_population, names=self.ps.names), name='DE')
             de.header['npop'] = self._tsa.de.n_pop
@@ -1141,11 +1220,8 @@ class ExoIris:
                                     log10_rho_bounds: float | tuple[float, float] = (-5, 0),
                                     log10_sigma_prior=None, log10_rho_prior=None,
                                     npop: int = 10, niter: int = 100):
-        if self._tsa.noise_model != 'fixed_gp':
-            raise ValueError("The noise model must be set to 'fixed_gp' before the hyperparameter optimization.")
-
-        if self._wa is None:
-            raise ValueError("The white light curves must be fit using 'fit_white()' before the hyperparameter optimization.")
+        if self._tsa.noise_model not in ('fixed_gp', 'free_gp'):
+            raise ValueError("The noise model must be set to 'fixed_gp' or 'free_gp' before the hyperparameter optimization.")
 
         if log10_rho_prior is not None:
             if isinstance(log10_rho_prior, Sequence):
@@ -1172,7 +1248,7 @@ class ExoIris:
 
             match log10_sigma_bounds:
                 case None:
-                    sb =  [log10_sigma_guess-1, log10_sigma_guess+1]
+                    sb = [log10_sigma_guess - 1, log10_sigma_guess + 1]
                 case _ if isinstance(log10_sigma_bounds, Sequence):
                     sb = log10_sigma_bounds
                 case _ if isinstance(log10_sigma_bounds, float):
@@ -1180,7 +1256,7 @@ class ExoIris:
 
             match log10_rho_bounds:
                 case None:
-                    rb =  [-5, -2]
+                    rb = [-5, -2]
                 case _ if isinstance(log10_rho_bounds, Sequence):
                     rb = log10_rho_bounds
                 case _ if isinstance(log10_rho_bounds, float):

exoiris/ldtkld.py

@@ -23,18 +23,18 @@ from numpy import sqrt, array, concatenate
 from pytransit import LDTkLD as PTLDTkLD
 from ldtk import BoxcarFilter
 
-from .tsdata import TSData, TSDataSet
+from .tsdata import TSData, TSDataGroup
 
 
 class LDTkLD(PTLDTkLD):
-    def __init__(self, data: TSDataSet | TSData,
+    def __init__(self, data: TSDataGroup | TSData,
                  teff: tuple[float, float],
                  logg: tuple[float, float],
                  metal: tuple[float, float],
                  cache: str | Path | None = None,
                  dataset: str = 'vis-lowres') -> None:
 
-        data = TSDataSet([data]) if isinstance(data, TSData) else data
+        data = TSDataGroup([data]) if isinstance(data, TSData) else data
         wl_edges = concatenate([array([d._wl_l_edges, d._wl_r_edges]) for d in data], axis=1).T
         filters = [BoxcarFilter(f"{0.5*(wla+wlb):08.5f}", wla*1e3, wlb*1e3) for wla, wlb in wl_edges]
         super().__init__(filters, teff, logg, metal, cache, dataset)