ExoIris 0.22.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

exoiris/exoiris.py

@@ -142,8 +142,11 @@ def load_model(fname: Path | str, name: str | None = None):
         # Read the priors.
         # ================
         priors = pickle.loads(codecs.decode(json.loads(hdul['PRIORS'].header['PRIORS']).encode(), "base64"))
-        a._tsa.ps = ParameterSet([pickle.loads(p) for p in priors])
-        a._tsa.ps.freeze()
+        for praw in priors:
+            p = pickle.loads(praw)
+            if p.name in a._tsa.ps.names:
+                a._tsa.set_prior(p.name, p.prior)
+
         if 'DE' in hdul:
             a._tsa._de_population = Table(hdul['DE'].data).to_pandas().values
             a._tsa._de_imin = hdul['DE'].header['IMIN']
@@ -281,9 +284,6 @@ class ExoIris:
         if parameter == 'radius ratios':
             for l in self._tsa.k_knots:
                 self.set_prior(f'k_{l:08.5f}', prior, *nargs)
-        elif parameter == 'baselines':
-            for par in self.ps[self._tsa._sl_baseline]:
-                self.set_prior(par.name, prior, *nargs)
         elif parameter == 'wn multipliers':
             for par in self.ps[self._tsa._sl_wnm]:
                 self.set_prior(par.name, prior, *nargs)
@@ -1192,7 +1192,7 @@ class ExoIris:
             wavelengths. The representation (radius ratio or depth) depends on the
             specified `kind`.
         """
-        if self.mcmc_chains is None:
+        if self.mcmc_chains is None and samples is None:
             raise ValueError("Cannot calculate posterior transmission spectrum before running the MCMC sampler.")
 
         if kind not in ('radius_ratio', 'depth'):
@@ -1206,8 +1206,11 @@ class ExoIris:
             wavelengths.sort()
 
         k_posteriors = zeros((samples.shape[0], wavelengths.size))
+        k_knots = self._tsa.k_knots.copy()
         for i, pv in enumerate(samples):
-            k_posteriors[i, :] = self._tsa._ip(wavelengths, self._tsa.k_knots, pv[self._tsa._sl_rratios])
+            if self._tsa.free_k_knot_ids is not None:
+                k_knots[self._tsa.free_k_knot_ids] = pv[self._tsa._sl_kloc]
+            k_posteriors[i, :] = self._tsa._ip(wavelengths, k_knots, pv[self._tsa._sl_rratios])
 
         if kind == 'radius_ratio':
             return wavelengths, k_posteriors

exoiris/tsdata.py

@@ -27,6 +27,8 @@ from matplotlib.figure import Figure
 from matplotlib.pyplot import subplots, setp
 from matplotlib.ticker import LinearLocator, FuncFormatter
 from numpy import (
+    any,
+    all,
     isfinite,
     where,
     all,
@@ -51,6 +53,9 @@ from numpy import (
     nanmedian,
     nanmean,
     unique,
+    ascontiguousarray,
+    vstack,
+    ones_like,
 )
 from pytransit.orbits import fold
 
@@ -68,7 +73,7 @@ class TSData:
                  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, epoch_group: int = 0, offset_group: int = 0,
-                 mask_nonfinite_errors: bool = True) -> None:
+                 mask_nonfinite_errors: bool = True, covs: ndarray | None = None) -> None:
         """
         Parameters
         ----------
@@ -99,7 +104,7 @@ class TSData:
         if transit_mask is not None and transit_mask.size != time.size:
             raise ValueError("The size of the out-of-transit mask array must match the size of the time array.")
 
-        if n_baseline < 1:
+        if n_baseline < 0:
             raise ValueError("n_baseline must be greater than zero.")
 
         if noise_group < 0:
@@ -126,7 +131,18 @@ class TSData:
             self.mask &= isfinite(errors)
         self.fluxes: ndarray = where(self.mask, fluxes, nan)
         self.errors: ndarray = where(self.mask, errors, nan)
+
+        if covs is not None:
+            self.covs: ndarray = covs
+        else:
+            ctime = self.time - self.time.mean()
+            self.covs = ascontiguousarray(vstack([ones(self.time.size)]+[ctime**i for i in range(1, n_baseline+1)]).T)
+            self.covs[:, 1:] /= self.covs[:, 1:].std(axis=0)
+
         self.transit_mask: ndarray = transit_mask if transit_mask is not None else ones(time.size, dtype=bool)
+        self._wlmask: ndarray = all(self.mask, 1)
+        self._wls_with_nan: ndarray = where(~self._wlmask)[0]
+
         self._ephemeris: Ephemeris | None = ephemeris
         self.n_baseline: int = n_baseline
         self.noise_group: int = noise_group
@@ -170,6 +186,7 @@ class TSData:
         time = pf.ImageHDU(self.time, name=f'time_{self.name}')
         wave = pf.ImageHDU(self.wavelength, name=f'wave_{self.name}')
         data = pf.ImageHDU(array([self.fluxes, self.errors]), name=f'data_{self.name}')
+        covs = pf.ImageHDU(self.covs, name=f'covs_{self.name}')
         ootm = pf.ImageHDU(self.transit_mask.astype(int), name=f'ootm_{self.name}')
         mask = pf.ImageHDU(self.mask.astype(int), name=f'mask_{self.name}')
         data.header['ngroup'] = self.noise_group
@@ -177,7 +194,7 @@ class TSData:
         data.header['epgroup'] = self.epoch_group
         data.header['offgroup'] = self.offset_group
         #TODO: export ephemeris
-        return pf.HDUList([time, wave, data, ootm, mask])
+        return pf.HDUList([time, wave, data, covs, ootm, mask])
 
     @staticmethod
     def import_fits(name: str, hdul: pf.HDUList) -> 'TSData':
@@ -200,6 +217,11 @@ class TSData:
         ootm = hdul[f'OOTM_{name}'].data.astype(bool)
         mask = hdul[f'MASK_{name}'].data.astype(bool)
 
+        try:
+            covs = hdul[f'COVS_{name}'].data.astype('d')
+        except KeyError:
+            covs = None
+
         try:
             noise_group = hdul[f'DATA_{name}'].header['NGROUP']
         except KeyError:
@@ -222,7 +244,8 @@ 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, epoch_group=ephemeris_group, offset_group=offset_group)
+                      n_baseline=n_baseline, mask=mask, epoch_group=ephemeris_group, offset_group=offset_group,
+                      covs=covs)
 
     def __repr__(self) -> str:
         return f"TSData Name:'{self.name}' [{self.wavelength[0]:.2f} - {self.wavelength[-1]:.2f}] nwl={self.nwl} npt={self.npt}"
@@ -301,6 +324,8 @@ class TSData:
         self.maxtm = self.time.max()
         if self._ephemeris is not None:
             self.mask_transit(ephemeris=self._ephemeris)
+        self._wlmask = all(self.mask, 1)
+        self._wls_with_nan = where(~self._wlmask)[0]
 
     def _update_data_mask(self) -> None:
         self.mask = isfinite(self.fluxes)
@@ -377,7 +402,8 @@ class TSData:
                    transit_mask=self.transit_mask[m],
                    ephemeris=self.ephemeris,
                    n_baseline=self.n_baseline,
-                   mask_nonfinite_errors=self.mask_nonfinite_errors)
+                   mask_nonfinite_errors=self.mask_nonfinite_errors,
+                   covs=self.covs[m])
         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],
@@ -387,7 +413,8 @@ class TSData:
                            transit_mask=self.transit_mask[m],
                            ephemeris=self.ephemeris,
                            n_baseline=self.n_baseline,
-                           mask_nonfinite_errors=self.mask_nonfinite_errors)
+                           mask_nonfinite_errors=self.mask_nonfinite_errors,
+                           covs=self.covs[m])
         return d
 
     def crop_wavelength(self, lmin: float, lmax: float, inplace: bool = True) -> 'TSData':
@@ -448,6 +475,7 @@ class TSData:
             self.transit_mask = self.transit_mask[m]
             self._tm_l_edges = self._tm_l_edges[m]
             self._tm_r_edges = self._tm_r_edges[m]
+            self.covs = self.covs[m]
             self._update()
             return self
         else:
@@ -463,7 +491,8 @@ class TSData:
                           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,
-                          mask_nonfinite_errors=self.mask_nonfinite_errors)
+                          mask_nonfinite_errors=self.mask_nonfinite_errors,
+                          covs=self.covs[m])
 
     # TODO: separate mask into bad data mask and outlier mask.
     def mask_outliers(self, sigma: float = 5.0) -> 'TSData':
@@ -486,6 +515,8 @@ class TSData:
         self.mask &= abs(self.fluxes - fm) / fe < sigma
         self.fluxes = where(self.mask, self.fluxes, nan)
         self.errors = where(self.mask, self.errors, nan)
+        self._wlmask = all(self.mask, 1)
+        self._wls_with_nan = where(~self._wlmask)[0]
         return self
 
     @deprecated("0.10", alternative="TSData.mask_outliers")
@@ -688,7 +719,8 @@ class TSData:
                           offset_group=self.offset_group,
                           transit_mask=self.transit_mask,
                           ephemeris=self.ephemeris,
-                          n_baseline=self.n_baseline)
+                          n_baseline=self.n_baseline,
+                          covs=self.covs)
 
     def bin_time(self, binning: Optional[Union[Binning, CompoundBinning]] = None,
                        nb: Optional[int] = None, bw: Optional[float] = None,
@@ -719,6 +751,7 @@ class TSData:
                 binning = Binning(self.time.min(), self.time.max(), nb=nb, bw=bw/(24*60*60) if bw is not None else None)
             bf, be = bin2d(self.fluxes.T, self.errors.T, self._tm_l_edges, self._tm_r_edges,
                            binning.bins, estimate_errors=estimate_errors)
+            bc, _ = bin2d(self.covs, ones_like(self.covs), self._tm_l_edges, self._tm_r_edges, binning.bins, False)
             d = TSData(binning.bins.mean(1), self.wavelength, bf.T, be.T,
                        wl_edges=(self._wl_l_edges, self._wl_r_edges),
                        tm_edges=(binning.bins[:,0], binning.bins[:,1]),
@@ -727,7 +760,8 @@ class TSData:
                        ephemeris=self.ephemeris,
                        n_baseline=self.n_baseline,
                        epoch_group=self.epoch_group,
-                       offset_group=self.offset_group)
+                       offset_group=self.offset_group,
+                       covs=bc)
             if self.ephemeris is not None:
                 d.mask_transit(ephemeris=self.ephemeris)
             return d
@@ -741,7 +775,7 @@ class TSDataGroup:
         self.wlmax: float = -inf
         self.tmin: float = inf
         self.tmax: float = -inf
-        self._noise_groups: array | None = None
+        self._noise_groups: ndarray | None = None
         for d in data:
             self._add_data(d)
 
@@ -782,7 +816,7 @@ class TSDataGroup:
         return [d.errors for d in self.data]
 
     @property
-    def noise_groups(self) -> ndarray[int]:
+    def noise_groups(self) -> ndarray[int] | None:
         """Array of noise groups."""
         return self._noise_groups
 
@@ -935,4 +969,4 @@ class TSDataGroup:
 
 
 class TSDataSet(TSDataGroup):
-    pass
+    pass

exoiris/tslpf.py

@@ -17,18 +17,35 @@
 from copy import deepcopy
 from typing import Optional, Literal
 
-from ldtk import BoxcarFilter, LDPSetCreator   # noqa
+from celerite2 import GaussianProcess as GP, terms
+from ldtk import BoxcarFilter, LDPSetCreator  # noqa
 from numba import njit, prange
-from numpy import zeros, log, pi, linspace, inf, atleast_2d, newaxis, clip, arctan2, ones, floor, sum, concatenate, \
-    sort, ndarray, zeros_like, array, tile, arange, squeeze, dstack, nan, diff, all
+from numpy import (
+    zeros,
+    log,
+    pi,
+    linspace,
+    inf,
+    atleast_2d,
+    newaxis,
+    clip,
+    arctan2,
+    sum,
+    concatenate,
+    sort,
+    ndarray,
+    array,
+    tile,
+    arange,
+    dstack,
+    diff,
+    ascontiguousarray,
+)
+from numpy.linalg import lstsq
 from numpy.random import default_rng
-from celerite2 import GaussianProcess as GP, terms
-
 from pytransit.lpf.logposteriorfunction import LogPosteriorFunction
-
-from pytransit.orbits import as_from_rhop, i_from_ba, fold, i_from_baew, d_from_pkaiews, epoch
+from pytransit.orbits import as_from_rhop, i_from_ba
 from pytransit.param import ParameterSet, UniformPrior as UP, NormalPrior as NP, GParameter
-from pytransit.stars import create_bt_settl_interpolator
 from scipy.interpolate import (
     pchip_interpolate,
     splrep,
@@ -37,10 +54,10 @@ from scipy.interpolate import (
     interp1d,
 )
 
-from .tsmodel import TransmissionSpectroscopyModel as TSModel
-from .tsdata import TSDataGroup
 from .ldtkld import LDTkLD
 from .spotmodel import SpotModel
+from .tsdata import TSDataGroup
+from .tsmodel import TransmissionSpectroscopyModel as TSModel
 
 NM_WHITE = 0
 NM_GP_FIXED = 1
@@ -49,6 +66,17 @@ NM_GP_FREE = 2
 noise_models = dict(white=NM_WHITE, fixed_gp=NM_GP_FIXED, free_gp=NM_GP_FREE)
 
 
+@njit
+def nlstsq(covs, res, mask, wlmask, with_nans):
+    nwl = res.shape[0]
+    nc = covs.shape[1]
+    x = zeros((nc, nwl))
+    x[:, wlmask] = lstsq(covs, ascontiguousarray(res[wlmask].T))[0]
+    for i in with_nans:
+        x[:, i] =  lstsq(covs[mask[i]], res[i, mask[i]])[0]
+    return x
+
+
 @njit(parallel=True, cache=False)
 def lnlike_normal(o, m, e, f, mask):
     nwl, nt = o.shape
@@ -244,7 +272,6 @@ class TSLPF(LogPosteriorFunction):
         self._init_p_noise()
         if self._nm == NM_GP_FREE:
             self._init_p_gp()
-        self._init_p_baseline()
         self._init_p_bias()
         self.ps.freeze()
 
@@ -399,23 +426,6 @@ class TSLPF(LogPosteriorFunction):
             self._start_gp = ps.blocks[-1].start
             self._sl_gp = ps.blocks[-1].slice
 
-    def _init_p_baseline(self):
-        ps = self.ps
-        self.n_baselines = self.data.n_baselines
-        self.baseline_knots = []
-        pp = []
-        for i, d in enumerate(self.data):
-            if d.n_baseline== 1:
-                self.baseline_knots.append([])
-                pp.append(GParameter(f'bl_{i:02d}_c', 'baseline constant', '', NP(1.0, 1e-6), (0, inf)))
-            elif d.n_baseline > 1:
-                knots = linspace(d.wavelength.min(), d.wavelength.max(), d.n_baseline)
-                self.baseline_knots.append(knots)
-                pp.extend([GParameter(f'bl_{i:02d}_{k:08.5f}', fr'baseline at {k:08.5f} $\mu$m', '', NP(1.0, 1e-6), (0, inf)) for k in knots])
-        ps.add_global_block('baseline_coefficients', pp)
-        self._start_baseline = ps.blocks[-1].start
-        self._sl_baseline = ps.blocks[-1].slice
-
     def _init_p_bias(self):
         ps = self.ps
         offset_groups = self.data.offset_groups
@@ -639,9 +649,8 @@ class TSLPF(LogPosteriorFunction):
             self.de = None
             self._de_population = pvpn
 
-    def _eval_k(self, pvp):
-        """
-        Evaluate the radius ratio model.
+    def _eval_k(self, pvp) -> list[ndarray]:
+        """Evaluate the radius ratio model.
 
         Parameters
         ----------
@@ -733,32 +742,26 @@ class TSLPF(LogPosteriorFunction):
                 fluxes[i] = biases + (1.0 - biases) * fluxes[i]
         return fluxes
 
-    def baseline_model(self, pv):
-        pv = atleast_2d(pv)[:, self._sl_baseline]
-        npv = pv.shape[0]
+    def baseline_model(self, mtransit):
+        npv = mtransit[0].shape[0]
         if self._baseline_models is None or self._baseline_models[0].shape[0] != npv:
-            self._baseline_models = [zeros((npv, d.nwl)) for d in self.data]
-        j = 0
+            self._baseline_models = [zeros(m.shape) for m in mtransit]
         for i, d in enumerate(self.data):
-            nbl = d.n_baseline
-            m = self._baseline_models[i]
-            if nbl == 1:
-                m[:, :] = pv[:, j][:, newaxis]
-            else:
-                for ipv in range(npv):
-                    m[ipv, :] = splev(d.wavelength, splrep(self.baseline_knots[i], pv[ipv, j:j+nbl], k=min(nbl-1, 3)))
-            j += nbl
+            for ipv in range(npv):
+                res = d.fluxes / mtransit[i][ipv]
+                coeffs = nlstsq(d.covs, res, d.mask, d._wlmask, d._wls_with_nan)
+                self._baseline_models[i][ipv, :, :] = (d.covs @ coeffs).T
         return self._baseline_models
 
     def flux_model(self, pv):
         transit_models = self.transit_model(pv)
-        baseline_models = self.baseline_model(pv)
+        baseline_models = self.baseline_model(transit_models)
         if self.spot_model is not None:
             self.spot_model.apply_spots(pv, transit_models)
             if self.spot_model.include_tlse:
                 self.spot_model.apply_tlse(pv, transit_models)
         for i in range(self.data.size):
-            transit_models[i][:, :, :] *= baseline_models[i][:, :, newaxis]
+            transit_models[i][:, :, :] *= baseline_models[i][:, :, :]
         return transit_models
 
     def create_pv_population(self, npop: int = 50) -> ndarray:

{exoiris-0.22.0.dist-info → exoiris-0.23.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ExoIris
-Version: 0.22.0
+Version: 0.23.0
 Summary: Easy and robust exoplanet transmission spectroscopy.
 Author-email: Hannu Parviainen <hannu@iac.es>
 License: GPLv3

{exoiris-0.22.0.dist-info → exoiris-0.23.0.dist-info}/RECORD

@@ -1,17 +1,17 @@
 exoiris/__init__.py,sha256=LU5jAE7_OVPLHFO0UAOGS0e0wuWV6rdSD0Qveet11K8,1147
 exoiris/binning.py,sha256=-Y9hdK0jZj8DOS82keaprneid2lZ4rCx-keWlKi0LP8,6455
 exoiris/ephemeris.py,sha256=dthBkJztT5yAP6VnnO7jGvxikboFUQBUGPUfBCFrA3w,1316
-exoiris/exoiris.py,sha256=U7bMoEovZ8lyerN4t7Rb3DA1jMnx0s3QVTo83_BTosQ,60863
+exoiris/exoiris.py,sha256=-N0fvkHRIUftkSNXaD1ioWeylG7vAb7gtTWWMQtNdDQ,60956
 exoiris/ldtkld.py,sha256=7H1r1xail3vSKdsNKorMTqivnRKU9WrOVH-uE4Ky2jM,3495
 exoiris/loglikelihood.py,sha256=lVm7kMyzy-xxP5VvZMV7kJY3KovbhlEz3u_6O7R6zgI,6309
 exoiris/spotmodel.py,sha256=9-DxvVzGzxf6AjQfrzZreyJB4Htw0gsIAD3nWl0tQMc,7160
-exoiris/tsdata.py,sha256=WqId5rfZR08pFZ83UZiyO39-QjX6WcB1GrUYolZsM-4,35323
-exoiris/tslpf.py,sha256=FYEAv2lNU4TQPVVuB7qIfQFefsru0tf_kZOnlQaRbS0,34096
+exoiris/tsdata.py,sha256=s3R6sByy7Ud26F3TdaRaXNQi-rYeJHnm7XHk2FLd1Go,36622
+exoiris/tslpf.py,sha256=iBagXEuTp2vNuH9jpGMXP7Ti5N9Hz0vzchaZpdD5Mr4,33343
 exoiris/tsmodel.py,sha256=6NaGY48fWHUT_7ti6Ao618PN-LgyoIhfQd8lZQqZ7hU,5160
 exoiris/util.py,sha256=uNv_c3Kuv1lml8MuDAuyElO4s3f1tRIQ1QMlLaI7Yak,5921
 exoiris/wlpf.py,sha256=g6h1cLk2-nKD8u_FzwXNVVGFK4dry8fBr0A70LA5gJw,6281
-exoiris-0.22.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-exoiris-0.22.0.dist-info/METADATA,sha256=Y5gNCFu6JH-lv65UPuY_Qts_HV9SgXKr1vRc1U0km48,5084
-exoiris-0.22.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-exoiris-0.22.0.dist-info/top_level.txt,sha256=EoNxT6c5mQDcM0f_LUQB-ETsYg03lNaV3o2L_Yc6-aE,8
-exoiris-0.22.0.dist-info/RECORD,,
+exoiris-0.23.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
+exoiris-0.23.0.dist-info/METADATA,sha256=fiRuio6EYjCval6pR1nV31KdL2wjvlgQxGMyOL743Hc,5084
+exoiris-0.23.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+exoiris-0.23.0.dist-info/top_level.txt,sha256=EoNxT6c5mQDcM0f_LUQB-ETsYg03lNaV3o2L_Yc6-aE,8
+exoiris-0.23.0.dist-info/RECORD,,