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

exoiris/exoiris.py

@@ -32,7 +32,8 @@ from emcee import EnsembleSampler
 from matplotlib.pyplot import subplots, setp, figure, Figure, Axes
 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)
+                   ceil, unique, zeros, cov)
+from numpy.typing import ArrayLike
 from numpy.random import normal
 from pytransit import UniformPrior, NormalPrior
 from pytransit.param import ParameterSet
@@ -44,6 +45,7 @@ from .ldtkld import LDTkLD
 from .tsdata import TSData, TSDataGroup
 from .tslpf import TSLPF
 from .wlpf import WhiteLPF
+from .loglikelihood import LogLikelihood
 
 
 def load_model(fname: Path | str, name: str | None = None):
@@ -131,6 +133,12 @@ def load_model(fname: Path | str, name: str | None = None):
             for i in range(hdr['NSPOTS']):
                 a.add_spot(hdr[f'SP{i+1:02d}_EG'])
 
+        # Read the free k knot indices if they exist.
+        # ==========================================
+        if 'N_FREE_K' in hdr and hdr['N_FREE_K'] > 0:
+            n_free_k = hdr['N_FREE_K']
+            a._tsa.set_free_k_knots([int(hdr[f'KK_IX_{i:03d}']) for i in range(n_free_k)])
+
         # Read the priors.
         # ================
         priors = pickle.loads(codecs.decode(json.loads(hdul['PRIORS'].header['PRIORS']).encode(), "base64"))
@@ -152,7 +160,7 @@ class ExoIris:
 
     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', 'nearest', 'linear'] = 'makima'):
+                 interpolation: Literal['nearest', 'linear', 'pchip', 'makima', 'bspline', 'bspline-quadratic'] = 'makima'):
         """
         Parameters
         ----------
@@ -209,6 +217,7 @@ class ExoIris:
         self._white_fluxes: None | list[ndarray] = None
         self._white_errors: None | list[ndarray] = None
         self._white_models: None | list[ndarray] = None
+        self.white_gp_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.
@@ -651,6 +660,9 @@ class ExoIris:
         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.
         """
+        if self.white_gp_models is None:
+            raise ValueError("White light curve GP predictions are not available. Run 'optimize_gp_hyperparameters' first.")
+
         ndata = self.data.size
 
         if axs is None:
@@ -869,13 +881,13 @@ class ExoIris:
 
         if result == 'fit':
             pv = self._tsa._de_population[self._tsa._de_imin]
-            ks = self._tsa._eval_k(pv[self._tsa._sl_rratios])
+            ks = self._tsa._eval_k(pv)
             ar = 1e2 * concatenate([squeeze(k) for k in ks]) ** 2
             ax.plot(wavelength[ix], ar[ix], c='k')
             ax.plot(self._tsa.k_knots, 1e2 * pv[self._tsa._sl_rratios] ** 2, 'k.')
         else:
             df = pd.DataFrame(self._tsa._mc_chains.reshape([-1, self._tsa.ndim]), columns=self._tsa.ps.names)
-            ks = self._tsa._eval_k(df.iloc[:, self._tsa._sl_rratios])
+            ks = self._tsa._eval_k(df.values)
             ar = 1e2 * concatenate(ks, axis=1) ** 2
             ax.fill_between(wavelength[ix], *percentile(ar[:, ix], [16, 84], axis=0), alpha=0.25)
             ax.plot(wavelength[ix], median(ar, 0)[ix], c='k')
@@ -891,7 +903,11 @@ class ExoIris:
             ax.set_xticks(xticks, labels=xticks)
         return ax.get_figure()
 
-    def plot_limb_darkening_parameters(self, result: Optional[str] = None, axs: Optional[tuple[Axes, Axes]] = None) -> None | Figure:
+    def plot_limb_darkening_parameters(
+        self,
+        result: None | Literal["fit", "mcmc"] = None,
+        axs: None | tuple[Axes, Axes] = None,
+    ) -> None | Figure:
         """Plot the limb darkening parameters.
 
         Parameters
@@ -920,56 +936,68 @@ class ExoIris:
         This method plots the limb darkening parameters for two-parameter limb darkening models. It supports only
         quadratic, quadratic-tri, power-2, and power-2-pm models.
         """
-        if not self._tsa.ldmodel in ('quadratic', 'quadratic-tri', 'power-2', 'power-2-pm'):
+        if not self._tsa.ldmodel in (
+            "quadratic",
+            "quadratic-tri",
+            "power-2",
+            "power-2-pm",
+        ):
             return None
 
         if axs is None:
-            fig, axs = subplots(1, 2, sharey='all', figsize=(13,4))
+            fig, axs = subplots(1, 2, sharey="all", figsize=(13, 4))
         else:
             fig = axs[0].get_figure()
 
         if result is None:
-            result = 'mcmc' if self._tsa.sampler is not None else 'fit'
-        if result not in ('fit', 'mcmc'):
+            result = "mcmc" if self._tsa.sampler is not None else "fit"
+        if result not in ("fit", "mcmc"):
             raise ValueError("Result must be either 'fit', 'mcmc', or None")
-        if result == 'mcmc' and self._tsa.sampler is None:
-            raise ValueError("Cannot plot posterior solution before running the MCMC sampler.")
+        if result == "mcmc" and not (
+            self._tsa.sampler is not None or self.mcmc_chains is not None
+        ):
+            raise ValueError(
+                "Cannot plot posterior solution before running the MCMC sampler."
+            )
 
         wavelength = concatenate(self.data.wavelengths)
         ix = argsort(wavelength)
 
-        if result == 'fit':
+        if result == "fit":
             pv = self._tsa._de_population[self._tsa._de_imin]
             ldc = squeeze(concatenate(self._tsa._eval_ldc(pv), axis=1))
-            axs[0].plot(self._tsa.ld_knots, pv[self._tsa._sl_ld][0::2], 'ok')
-            axs[0].plot(wavelength[ix], ldc[:,0][ix])
-            axs[1].plot(self._tsa.ld_knots, pv[self._tsa._sl_ld][1::2], 'ok')
-            axs[1].plot(wavelength[ix], ldc[:,1][ix])
+            axs[0].plot(self._tsa.ld_knots, pv[self._tsa._sl_ld][0::2], "ok")
+            axs[0].plot(wavelength[ix], ldc[:, 0][ix])
+            axs[1].plot(self._tsa.ld_knots, pv[self._tsa._sl_ld][1::2], "ok")
+            axs[1].plot(wavelength[ix], ldc[:, 1][ix])
         else:
-            pvp = self._tsa._mc_chains.reshape([-1, self._tsa.ndim])
-            ldc = pvp[:,self._tsa._sl_ld]
+            if self._tsa.sampler is not None:
+                pvp = self._tsa._mc_chains.reshape([-1, self._tsa.ndim])
+            else:
+                pvp = self.mcmc_chains.reshape([-1, self._tsa.ndim])
+            ldc = pvp[:, self._tsa._sl_ld]
 
-            ld1m = median(ldc[:,::2], 0)
-            ld1e = ldc[:,::2].std(0)
-            ld2m = median(ldc[:,1::2], 0)
-            ld2e = ldc[:,1::2].std(0)
+            ld1m = median(ldc[:, ::2], 0)
+            ld1e = ldc[:, ::2].std(0)
+            ld2m = median(ldc[:, 1::2], 0)
+            ld2e = ldc[:, 1::2].std(0)
 
             ldc = concatenate(self._tsa._eval_ldc(pvp), axis=1)
-            ld1p = percentile(ldc[:,:,0], [50, 16, 84], axis=0)
-            ld2p = percentile(ldc[:,:,1], [50, 16, 84], axis=0)
+            ld1p = percentile(ldc[:, :, 0], [50, 16, 84], axis=0)
+            ld2p = percentile(ldc[:, :, 1], [50, 16, 84], axis=0)
 
             axs[0].fill_between(wavelength[ix], ld1p[1, ix], ld1p[2, ix], alpha=0.5)
-            axs[0].plot(wavelength[ix], ld1p[0][ix], 'k')
+            axs[0].plot(wavelength[ix], ld1p[0][ix], "k")
             axs[1].fill_between(wavelength[ix], ld2p[1, ix], ld2p[2, ix], alpha=0.5)
-            axs[1].plot(wavelength[ix], ld2p[0][ix], 'k')
+            axs[1].plot(wavelength[ix], ld2p[0][ix], "k")
 
-            axs[0].errorbar(self._tsa.ld_knots, ld1m, ld1e, fmt='ok')
-            axs[1].errorbar(self._tsa.ld_knots, ld2m, ld2e, fmt='ok')
+            axs[0].errorbar(self._tsa.ld_knots, ld1m, ld1e, fmt="ok")
+            axs[1].errorbar(self._tsa.ld_knots, ld2m, ld2e, fmt="ok")
 
         ldp = full((self.nldp, 2, 2), nan)
         for i in range(self.nldp):
             for j in range(2):
-                p = self.ps[self._tsa._sl_ld][i*2+j].prior
+                p = self.ps[self._tsa._sl_ld][i * 2 + j].prior
                 if isinstance(p, UniformPrior):
                     ldp[i, j, 0] = p.a
                     ldp[i, j, 1] = p.b
@@ -979,11 +1007,15 @@ class ExoIris:
 
         for i in range(2):
             for j in range(2):
-                axs[i].plot(self._tsa.ld_knots, ldp[:, i, j], ':', c='C0')
-
-        setp(axs, xlim=(wavelength.min(), wavelength.max()), xlabel=r'Wavelength [$\mu$m]')
-        setp(axs[0], ylabel='Limb darkening coefficient 1')
-        setp(axs[1], ylabel='Limb darkening coefficient 2')
+                axs[i].plot(self._tsa.ld_knots, ldp[:, i, j], ":", c="C0")
+
+        setp(
+            axs,
+            xlim=(wavelength.min(), wavelength.max()),
+            xlabel=r"Wavelength [$\mu$m]",
+        )
+        setp(axs[0], ylabel="Limb darkening coefficient 1")
+        setp(axs[1], ylabel="Limb darkening coefficient 2")
         return fig
 
     def plot_residuals(self, result: Optional[str] = None, ax: None | Axes | Sequence[Axes] = None,
@@ -1107,8 +1139,8 @@ class ExoIris:
         return fig
 
     @property
-    def transmission_spectrum(self) -> Table:
-        """Get the posterior transmission spectrum as a Pandas DataFrame.
+    def transmission_spectrum_table(self) -> Table:
+        """Get the posterior transmission spectrum as an Astropy Table.
 
         Raises
         ------
@@ -1120,7 +1152,7 @@ class ExoIris:
 
         pvp = self.posterior_samples
         wls = concatenate(self.data.wavelengths)
-        ks = concatenate(self._tsa._eval_k(pvp.values[:, self._tsa._sl_rratios]), axis=1)
+        ks = concatenate(self._tsa._eval_k(pvp.values), axis=1)
         ar = ks**2
         ix = argsort(wls)
         return Table(data=[wls[ix]*u.micrometer,
@@ -1128,6 +1160,101 @@ class ExoIris:
                            median(ar, 0)[ix], ar.std(0)[ix]],
                      names = ['wavelength', 'radius_ratio', 'radius_ratio_e', 'area_ratio', 'area_ratio_e'])
 
+    def transmission_spectrum_samples(self, wavelengths: ndarray | None = None,
+                                      kind: Literal['radius_ratio', 'depth'] = 'depth',
+                                      samples: ndarray | None = None) -> tuple[ndarray, ndarray]:
+        """Calculate posterior transmission spectrum samples.
+
+        This method computes the posterior samples of the transmission spectrum,
+        either as radius ratios or as transit depths, depending on the specified
+        kind. It interpolates the data for given wavelengths or uses the
+        instrumental wavelength grid if none is provided. Requires that MCMC
+        sampling has been performed prior to calling this method.
+
+        Parameters
+        ----------
+        wavelengths
+            The array of wavelengths at which the spectrum should be sampled.
+            If None, the default wavelength grid defined by the instrumental data
+            will be used.
+        kind
+            Specifies the desired representation of the transmission spectrum.
+            'radius_ratio' returns the spectrum in radius ratio units, while
+            'depth' returns the spectrum in transit depth units. Default is 'depth'.
+        samples
+            Array of posterior samples to use for calculation. If None,
+            the method will use previously stored posterior samples.
+
+        Returns
+        -------
+        ndarray
+            Array containing the transmission spectrum samples for the specified
+            wavelengths. The representation (radius ratio or depth) depends on the
+            specified `kind`.
+        """
+        if self.mcmc_chains is None:
+            raise ValueError("Cannot calculate posterior transmission spectrum before running the MCMC sampler.")
+
+        if kind not in ('radius_ratio', 'depth'):
+            raise ValueError("Invalid value for `kind`. Must be either 'radius_ratio' or 'depth'.")
+
+        if samples is None:
+            samples = self.posterior_samples.values
+
+        if wavelengths is None:
+            wavelengths = concatenate(self.data.wavelengths)
+            wavelengths.sort()
+
+        k_posteriors = zeros((samples.shape[0], wavelengths.size))
+        for i, pv in enumerate(samples):
+            k_posteriors[i, :] = self._tsa._ip(wavelengths, self._tsa.k_knots, pv[self._tsa._sl_rratios])
+
+        if kind == 'radius_ratio':
+            return wavelengths, k_posteriors
+        else:
+            return wavelengths, k_posteriors**2
+
+    def transmission_spectrum(self, wavelengths: ndarray | None = None, kind: Literal['radius_ratio', 'depth'] = 'depth', samples: ndarray | None = None, return_cov: bool = True) -> tuple[ndarray, ndarray]:
+        """Compute the transmission spectrum.
+
+        This method calculates the mean transmission spectrum values and the covariance matrix
+        (or standard deviations) for the given parameter set. The mean represents the average
+        transmission spectrum, and the covariance provides information on the uncertainties and
+        correlations between wavelengths or samples.
+
+        Parameters
+        ----------
+        wavelengths
+            Array of wavelength values at which to calculate the transmission spectrum.
+            If None, the default grid will be used.
+        kind
+            Specifies the method to represent the spectrum. 'radius_ratio' computes the
+            spectrum in terms of the planet-to-star radius ratio, while 'depth' computes
+            the spectrum in terms of transit depth.
+        samples
+            Array of samples used to compute the spectrum uncertainties. If None, previously
+            stored samples will be utilized.
+
+        return_cov : bool, optional
+            Indicates whether to return the covariance matrix of the computed transmission
+            spectrum. If True, the covariance matrix is returned along with the mean spectrum.
+            If False, the standard deviation of the spectrum is returned.
+
+        Returns
+        -------
+        tuple[ndarray, ndarray]
+            A tuple containing two arrays:
+            - The mean transmission spectrum.
+            - The covariance matrix of the spectrum (if `return_cov` is True), or the
+              standard deviation (if `return_cov` is False).
+        """
+        sp_samples = self.transmission_spectrum_samples(wavelengths, kind, samples)[1]
+        mean = sp_samples.mean(0)
+        if return_cov:
+            return mean, cov(sp_samples, rowvar=False)
+        else:
+            return mean, sp_samples.std(0)
+
     def save(self, overwrite: bool = False) -> None:
         """Save the ExoIris analysis to a FITS file.
 
@@ -1145,6 +1272,13 @@ class ExoIris:
         pri.header['interp'] = self._tsa.interpolation
         pri.header['noise'] = self._tsa.noise_model
 
+        if self._tsa.free_k_knot_ids is None:
+            pri.header['n_free_k'] = 0
+        else:
+            pri.header['n_free_k'] = len(self._tsa.free_k_knot_ids)
+            for i, ix in enumerate(self._tsa.free_k_knot_ids):
+                pri.header[f'kk_ix_{i:03d}'] = ix
+
         # Priors
         # ======
         pr = pf.ImageHDU(name='priors')
@@ -1228,6 +1362,33 @@ class ExoIris:
 
         hdul.writeto(f"{self.name}.fits", overwrite=True)
 
+    def create_loglikelihood_function(self, wavelengths: ndarray, kind: Literal['radius_ratio', 'depth'] = 'depth',
+                                      method: Literal['svd', 'randomized_svd', 'eigh'] = 'svd',
+                                      n_max_samples: int = 10000) -> LogLikelihood:
+        """Create a reduced-rank Gaussian log-likelihood function for retrieval.
+
+        Parameters
+        ----------
+        wavelengths
+            The wavelength grid used in the theoretical transmission spectra.
+
+        kind
+            The transmission spectrum type. Can be either 'radius_ratio' or 'depth'.
+
+        Returns
+        -------
+        LogLikelihood
+            An instance of LogLikelihood for analyzing the consistency of the model
+            with the provided wavelengths and chosen log-likelihood kind.
+        """
+        if self.mcmc_chains is None:
+            raise ValueError("Cannot create log-likelihood function before running the MCMC sampler.")
+        return LogLikelihood(wavelengths,
+                             self.transmission_spectrum_samples(wavelengths, kind)[1],
+                             method=method,
+                             n_max_samples=n_max_samples,
+                             nk=self.nk)
+
     def create_initial_population(self, n: int, source: str, add_noise: bool = True) -> ndarray:
         """Create an initial parameter vector population for the DE optimisation.
 

exoiris/loglikelihood.py

@@ -0,0 +1,139 @@
+#  ExoIris: fast, flexible, and easy exoplanet transmission spectroscopy in Python.
+#  Copyright (C) 2025 Hannu Parviainen
+#
+#  This program is free software: you can redistribute it and/or modify
+#  it under the terms of the GNU General Public License as published by
+#  the Free Software Foundation, either version 3 of the License, or
+#  (at your option) any later version.
+#
+#  This program is distributed in the hope that it will be useful,
+#  but WITHOUT ANY WARRANTY; without even the implied warranty of
+#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#  GNU General Public License for more details.
+#
+#  You should have received a copy of the GNU General Public License
+#  along with this program.  If not, see <https://www.gnu.org/licenses/>.
+from typing import Literal
+
+from numpy import full, cov, sum, ndarray, log, pi, asarray
+from numpy.linalg import eigh, svd
+from sklearn.utils.extmath import randomized_svd
+
+
+class LogLikelihood:
+    def __init__(self, wavelength: ndarray, spectra: None | ndarray = None, spmean: None | ndarray = None,
+                 spcov: None | ndarray = None, eps: float = 1e-10, method: Literal['svd', 'randomized_svd', 'eigh'] = 'svd',
+                 n_max_samples: int = 10000, nk: int | None = None):
+        """Reduced-rank Normal log-likelihood.
+
+        This class constructs a statistically robust log-likelihood function for
+        comparing a theoretical transmission spectrum to the posterior distribution
+        inferred by ExoIris.
+
+        Because the posterior samples are generated from a spline with $K$ knots
+        but evaluated on $M$ wavelengths ($M \gg K$), the empirical covariance
+        matrix is singular or strongly ill-conditioned. This class solves the
+        rank-deficiency problem by projecting the model into the principal
+        subspace of the posterior (Karhunen-Loève compression).
+
+        Parameters
+        ----------
+        wavelength
+            The wavelength grid with a shape (M,) on which the posterior samples and theoretical
+            spectra are evaluated.
+        spectra
+            The posterior spectrum samples with shape (N_samples, M_wavelengths).
+            If provided, ``spmean`` and ``spcov`` are computed automatically.
+            Mutually exclusive with ``spmean`` and ``spcov``.
+        spmean
+            The pre-computed mean spectrum with shape (M,). Must be provided
+            along with ``spcov`` if ``spectra`` is None.
+        spcov
+            The pre-computed covariance matrix with shape (M, M). Must be provided
+            along with ``spmean`` if ``spectra`` is None.
+        eps
+            Relative tolerance factor used to determine which eigenvalues of
+            the covariance matrix are considered significant. Eigenvalues smaller
+            than ``eps * max_eigenvalue`` are discarded. Default is ``1e-10``.
+
+        Notes
+        -----
+        This implementation follows the "Signal-to-Noise Eigenmode" formalism
+        described by Tegmark et al. (1997) for analyzing rank-deficient
+        cosmological datasets.
+
+        The log-likelihood is evaluated as:
+
+        .. math:: \ln \mathcal{L} = -\frac{1}{2} \left[ \sum_{i=1}^{K} \frac{p_i^2}{\lambda_i} + \sum_{i=1}^{K} \ln(\lambda_i) + K \ln(2\pi) \right]
+
+        where $\lambda_i$ are the significant eigenvalues of the covariance
+        matrix, and $p_i$ are the projections of the model residuals onto the
+        corresponding eigenvectors (principal components).
+
+        References
+        ----------
+        Tegmark, M., Taylor, A. N., & Heavens, A. F. (1997). Karhunen-Loève
+        eigenvalue problems in cosmology: how should we tackle large data sets?
+        *The Astrophysical Journal*, 480(1), 22.
+        """
+        self.wavelength = wavelength
+        self.eps = eps
+
+        if spectra is not None and (spmean is not None or spcov is not None):
+            raise ValueError("Cannot specify both `spectra` and `spmean` and `spcov`.")
+
+        if spectra is None and (spmean is None or spcov is None):
+            raise ValueError("Must specify either `spectra` or both `spmean` and `spcov`.")
+
+        if spectra is not None:
+            spectra = spectra[:n_max_samples, :]
+            self.spmean = spectra.mean(axis=0)
+
+        if method == 'svd':
+            _, sigma, evecs = svd(spectra - spectra.mean(0), full_matrices=False)
+            evals = (sigma**2) / (spectra.shape[0] - 1)
+            evecs = evecs.T
+        elif method == 'randomized_svd':
+            if nk is None:
+                raise ValueError("Must specify `nk` when using `method='randomized_svd'`.")
+            _, sigma, evecs = randomized_svd(spectra - spectra.mean(0),  n_components=nk, n_iter=5, random_state=0)
+            evals = (sigma ** 2) / (spectra.shape[0] - 1)
+            evecs = evecs.T
+        elif method == 'eigh' or (spmean is not None and spcov is not None):
+            if spectra is not None:
+                self.spcov = cov(spectra, rowvar=False)
+            else:
+                self.spmean = spmean
+                self.spcov = spcov
+            evals, evecs = eigh(self.spcov)
+
+        keep = evals > eps * evals.max()
+        self.eigenvalues, self.eigenvectors = evals[keep], evecs[:, keep]
+        self.log_det = sum(log(self.eigenvalues))
+        self.log_twopi = self.eigenvalues.size * log(2*pi)
+
+    def __call__(self, model: ndarray | float) -> ndarray:
+        """Evaluate the log-likelihood of a model spectrum.
+
+        Parameters
+        ----------
+        model : float or ndarray
+            The theoretical model spectrum. If a float is provided, it is
+            broadcast to a flat spectrum. If an array, it must match the
+            wavelength grid size used during initialization.
+
+        Returns
+        -------
+        float
+            The natural log-likelihood $\ln \mathcal{L}$.
+        """
+        if isinstance(model, float):
+            model = full(self.wavelength.size, model)
+        else:
+            model = asarray(model)
+
+        # Project the residuals onto the eigenvectors (Basis Rotation)
+        # and Compute the Mahalanobis Distance (Chi-Squared in Subspace).
+        p = (self.spmean - model) @ self.eigenvectors
+        chisq = sum(p**2 / self.eigenvalues)
+        return -0.5 * (chisq + self.log_det + self.log_twopi)

exoiris/tslpf.py

@@ -20,7 +20,7 @@ from typing import Optional, Literal
 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
+    sort, ndarray, zeros_like, array, tile, arange, squeeze, dstack, nan, diff, all
 from numpy.random import default_rng
 from celerite2 import GaussianProcess as GP, terms
 
@@ -35,7 +35,6 @@ from scipy.interpolate import (
     splev,
     Akima1DInterpolator,
     interp1d,
-    FloaterHormannInterpolator,
 )
 
 from .tsmodel import TransmissionSpectroscopyModel as TSModel
@@ -86,6 +85,10 @@ def ip_bspline(x, xk, yk):
     return splev(x, splrep(xk, yk))
 
 
+def ip_bspline_quadratic(x, xk, yk):
+    return splev(x, splrep(xk, yk, k=2))
+
+
 def ip_makima(x, xk, yk):
     return Akima1DInterpolator(xk, yk, method='makima', extrapolate=True)(x)
 
@@ -102,6 +105,13 @@ def add_knots(x_new, x_old):
     return sort(concatenate([x_new, x_old]))
 
 
+interpolator_choices = ("bspline", "pchip", "makima", "nearest", "linear", "bspline-quadratic")
+
+
+interpolators = {'bspline': ip_bspline, 'bspline-quadratic': ip_bspline_quadratic, 'pchip': ip_pchip,
+                 'makima': ip_makima, 'nearest': ip_nearest, 'linear': ip_linear}
+
+
 def clean_knots(knots, min_distance, lmin=0, lmax=inf):
     """Clean the knot table by replacing groups of adjacent knots with a single knot at the group mean.
 
@@ -143,19 +153,21 @@ def clean_knots(knots, min_distance, lmin=0, lmax=inf):
 class TSLPF(LogPosteriorFunction):
     def __init__(self, runner, name: str, ldmodel, data: TSDataGroup, nk: int = 50, nldc: int = 10, nthreads: int = 1,
                  tmpars = None, noise_model: Literal["white", "fixed_gp", "free_gp"] = 'white',
-                 interpolation: Literal['bspline', 'pchip', 'makima', 'nearest', 'linear'] = 'makima'):
+                 interpolation: Literal['nearest', 'linear', 'pchip', 'makima', 'bspline', 'bspline-quadratic'] = 'makima'):
         super().__init__(name)
         self._runner = runner
         self._original_data: TSDataGroup | None = None
         self.data: TSDataGroup | None = None
         self.npb: list[int] | None= None
         self.npt: list[int] | None = None
-        self.ndim: int | None = None
         self._baseline_models: list[ndarray] | None = None
         self.interpolation: str = interpolation
 
-        self._ip = {'bspline': ip_bspline, 'pchip': ip_pchip, 'makima': ip_makima,
-                    'nearest': ip_nearest, 'linear': ip_linear}[interpolation]
+        if interpolation not in interpolator_choices:
+            raise ValueError(f'interpolation must be one of {interpolator_choices}')
+
+        self._ip = interpolators[interpolation]
+        self._ip_ld = interpolators['bspline']
 
         self._gp: Optional[list[GP]] = None
         self._gp_time: Optional[list[ndarray]] = None
@@ -179,6 +191,7 @@ class TSLPF(LogPosteriorFunction):
         self.nk = nk
 
         self.k_knots = linspace(data.wlmin, data.wlmax, self.nk)
+        self.free_k_knot_ids = None
 
         if isinstance(ldmodel, LDTkLD):
             self.ld_knots = array([])
@@ -209,6 +222,10 @@ class TSLPF(LogPosteriorFunction):
     def errors(self) -> list[ndarray]:
         return self.data.errors
 
+    @property
+    def ndim(self) -> int:
+        return len(self.ps)
+
     def set_data(self, data: TSDataGroup):
         self._original_data = deepcopy(data)
         self.data = data
@@ -230,7 +247,6 @@ class TSLPF(LogPosteriorFunction):
         self._init_p_baseline()
         self._init_p_bias()
         self.ps.freeze()
-        self.ndim = len(self.ps)
 
     def initialize_spots(self, tstar: float, wlref: float, include_tlse: bool = True) -> None:
         self.spot_model = SpotModel(self, tstar, wlref, include_tlse)
@@ -523,6 +539,56 @@ class TSLPF(LogPosteriorFunction):
             self._mc_chains = fmcn.reshape([mco.shape[0], mco.shape[1], ndn])
             self.sampler = None
 
+    def set_free_k_knots(self, ids):
+        self.free_k_knot_ids = ids
+
+        # Remove existing parameter block if one exists
+        block_names = [b.name for b in self.ps.blocks]
+        try:
+            bid = block_names.index('free_k_knot_locations')
+            del self.ps[self.ps.blocks[bid].slice]
+            del self.ps.blocks[bid]
+        except ValueError:
+            pass
+
+        # Calculate minimum distances between knots
+        min_distances = zeros(self.nk)
+        min_distances[0] = self.k_knots[1] - self.k_knots[0]
+        min_distances[self.nk-1] = self.k_knots[self.nk-1] - self.k_knots[self.nk-2]
+        for i in range(1, self.nk-1):
+            for j in range(i):
+                min_distances[i] = min(self.k_knots[i] - self.k_knots[i-1], self.k_knots[i+1] - self.k_knots[i])
+
+        # Create new parameter block
+        ps = []
+        for kid in ids:
+            sigma = min_distances[kid]/6 if (kid+1 in ids or kid-1 in ids) else min_distances[kid]/4
+            ps.append(GParameter(f'kl_{kid:04d}', f'k knot {kid} location', 'um', NP(self.k_knots[kid], sigma), [0, inf]))
+        self.ps.thaw()
+        self.ps.add_global_block('free_k_knot_locations', ps)
+        self.ps.freeze()
+        self._start_kloc = self.ps.blocks[-1].start
+        self._sl_kloc = self.ps.blocks[-1].slice
+
+        try:
+            pid = [p.__name__ for p in self._additional_log_priors].index('k_knot_order_priors')
+            del self._additional_log_priors[pid]
+        except ValueError:
+            pass
+
+        # Add a prior on the order of the knots
+        def k_knot_order_prior(pv):
+            pv = atleast_2d(pv)
+            logp = zeros(pv.shape[0])
+            k_knots = self.k_knots.copy()
+            for i in range(pv.shape[0]):
+                k_knots[self.free_k_knot_ids] = pv[i, self._sl_kloc]
+                original_separations = diff(self.k_knots)
+                current_separations = diff(k_knots)
+                logp[i] = 1e2*(clip(current_separations / original_separations / 0.25, -inf, 1.0) - 1.).sum()
+            return logp
+        self._additional_log_priors.append(k_knot_order_prior)
+
 
     def add_ld_knots(self, knot_wavelengths) -> None:
         """Add limb darkening knots to the model.
@@ -590,9 +656,12 @@ class TSLPF(LogPosteriorFunction):
         """
         pvp = atleast_2d(pvp)
         ks = [zeros((pvp.shape[0], npb)) for npb in self.npb]
+        k_knots = self.k_knots.copy()
         for ids in range(self.data.size):
             for ipv in range(pvp.shape[0]):
-                ks[ids][ipv,:] =  self._ip(self.wavelengths[ids], self.k_knots, pvp[ipv])
+                if self.free_k_knot_ids is not None:
+                    k_knots[self.free_k_knot_ids] = pvp[ipv, self._sl_kloc]
+                ks[ids][ipv,:] =  self._ip(self.wavelengths[ids], k_knots, pvp[ipv, self._sl_rratios])
         return ks
 
     def _eval_ldc(self, pvp):
@@ -608,8 +677,8 @@ class TSLPF(LogPosteriorFunction):
             ldp = [zeros((pvp.shape[0], npb, 2)) for npb in self.npb]
             for ids in range(self.data.size):
                 for ipv in range(pvp.shape[0]):
-                    ldp[ids][ipv, :, 0] = ip_bspline(self.wavelengths[ids], self.ld_knots, ldk[ipv, :, 0])
-                    ldp[ids][ipv, :, 1] = ip_bspline(self.wavelengths[ids], self.ld_knots, ldk[ipv, :, 1])
+                    ldp[ids][ipv, :, 0] = self._ip_ld(self.wavelengths[ids], self.ld_knots, ldk[ipv, :, 0])
+                    ldp[ids][ipv, :, 1] = self._ip_ld(self.wavelengths[ids], self.ld_knots, ldk[ipv, :, 1])
             return ldp
 
     def transit_model(self, pv, copy=True):
@@ -640,7 +709,7 @@ class TSLPF(LogPosteriorFunction):
         pv = atleast_2d(pv)
         ldp = self._eval_ldc(pv)
         t0s = pv[:, self._sl_tcs]
-        k = self._eval_k(pv[:, self._sl_rratios])
+        k = self._eval_k(pv)
         p = pv[:, 1]
         aor = as_from_rhop(pv[:, 0], p)
         inc = i_from_ba(pv[:, 2], aor)

exoiris/util.py

@@ -15,8 +15,10 @@
 #  along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 from numba import njit
-from numpy import zeros, sum, sqrt, linspace, vstack, concatenate, floor, dot, ndarray, nan
-
+from numpy import (zeros, sum, sqrt, linspace, vstack, concatenate, floor, dot, ndarray, nan, asarray, tile)
+from numpy._typing import ArrayLike
+from pytransit import TSModel
+from pytransit.orbits import i_from_ba
 
 @njit
 def bin2d(v, e, el, er, bins, estimate_errors: bool = False) -> tuple[ndarray, ndarray]:
@@ -116,4 +118,45 @@ def create_binning(ranges, bwidths):
         n = int(floor((r[1] - r[0]) / w))
         e = linspace(*r, num=n)
         bins.append(vstack([e[:-1], e[1:]]).T)
-    return concatenate(bins)
+    return concatenate(bins)
+
+
+def create_mock_model(ks: ArrayLike, times: ArrayLike = None, ldc: ArrayLike = None, t0: float = 0.0, p: float =2.0, a: float =8.0, b: float =0.0) -> ndarray:
+    """Create a mock transmission spectrum observation using given parameters.
+
+    Parameters
+    ----------
+    ks
+        Array of radius ratios, one radius ratio per wavelength.
+    times
+        Array of time values to set the data points. If None, defaults to a
+        linspace of 500 points in the range [-0.1, 0.1].
+    ldc
+        Array representing the limb darkening coefficients. If None, defaults to
+        a tile of [0.4, 0.4] for each wavelength element.
+    t0
+        Transit center.
+    p
+        Orbital period.
+    a
+        Semi-major axis.
+    b
+        Impact parameter.
+
+    Returns
+    -------
+    ndarray
+        Mock spectrophotometric light curves.
+
+    """
+    ks = asarray(ks)
+    if times is None:
+        times = linspace(-0.1, 0.1, 500)
+    if ldc is None:
+        ldc = tile([0.4, 0.4], (1, ks.size, 1))
+    inc = i_from_ba(b, a)
+
+    m1 = TSModel('power-2', ng=100, nzin=50, nzlimb=50)
+    m1.set_data(times)
+    f1 = m1.evaluate(ks, ldc, t0, p, a, inc)[0]
+    return f1

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ExoIris
-Version: 0.20.0
+Version: 0.22.0
 Summary: Easy and robust exoplanet transmission spectroscopy.
 Author-email: Hannu Parviainen <hannu@iac.es>
 License: GPLv3
@@ -29,6 +29,7 @@ Requires-Dist: xarray
 Requires-Dist: seaborn
 Requires-Dist: astropy
 Requires-Dist: uncertainties
+Requires-Dist: scikit-learn
 Dynamic: license-file
 
 # ExoIris: Fast and Flexible Transmission Spectroscopy in Python

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

@@ -1,16 +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=POqkyc9hM5YmARDRtVN6k0l6N63hk6g_ecOR8Y4Z_2o,53838
+exoiris/exoiris.py,sha256=U7bMoEovZ8lyerN4t7Rb3DA1jMnx0s3QVTo83_BTosQ,60863
 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=yFyMjvx8P9t55wYGGfoEZ8u3Yw1H7T4Ot-dsnT2oQEw,31316
+exoiris/tslpf.py,sha256=FYEAv2lNU4TQPVVuB7qIfQFefsru0tf_kZOnlQaRbS0,34096
 exoiris/tsmodel.py,sha256=6NaGY48fWHUT_7ti6Ao618PN-LgyoIhfQd8lZQqZ7hU,5160
-exoiris/util.py,sha256=5PynwYYHRrzyXJHskBtp2J-pcM59zsA1_VtDxencQm4,4630
+exoiris/util.py,sha256=uNv_c3Kuv1lml8MuDAuyElO4s3f1tRIQ1QMlLaI7Yak,5921
 exoiris/wlpf.py,sha256=g6h1cLk2-nKD8u_FzwXNVVGFK4dry8fBr0A70LA5gJw,6281
-exoiris-0.20.0.dist-info/licenses/LICENSE,sha256=OXLcl0T2SZ8Pmy2_dmlvKuetivmyPd5m1q-Gyd-zaYY,35149
-exoiris-0.20.0.dist-info/METADATA,sha256=LfJuP-1bFOckme0rDdzOcohvnUeuvspSNQ_m9_yNrbc,5056
-exoiris-0.20.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-exoiris-0.20.0.dist-info/top_level.txt,sha256=EoNxT6c5mQDcM0f_LUQB-ETsYg03lNaV3o2L_Yc6-aE,8
-exoiris-0.20.0.dist-info/RECORD,,
+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,,