Added dataset-specific GP hyperparameters support and refined data handling

- Enhanced `set_gp_hyperparameters` to allow dataset-specific hyperparameters via `idata` parameter.
- Improved flux and error handling in `wlpf` by ensuring finite data treatment.
- Fixed initialization logic for `optimize` method in `wlpf`.
- Adjusted setup order for data and noise models in `tslpf` for consistency.
- Streamlined GP initialization logic by enhancing condition handling.

Author: hpparvi

exoiris/exoiris.py

@@ -303,7 +303,7 @@ def set_ldtk_prior(self,
         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
@@ -312,8 +312,10 @@ def set_gp_hyperparameters(self, sigma: float, rho: float) -> None:
             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.

exoiris/tslpf.py

@@ -142,13 +142,11 @@ def __init__(self, name: str, ldmodel, data: TSDataSet, nk: int = 50, nldc: int
         self._gp_time: Optional[list[ndarray]] = None
         self._gp_flux: Optional[list[ndarray]] = None
 
-        self.set_noise_model(noise_model)
-
-        self.ldmodel = ldmodel
-
         self.tms = [TSModel(ldmodel, nthreads=nthreads, **(tmpars or {})) for i in range(len(data))]
         self.set_data(data)
+        self.set_noise_model(noise_model)
 
+        self.ldmodel = ldmodel
         if isinstance(ldmodel, LDTkLD):
             for tm in self.tms:
                 tm.ldmodel = None
@@ -196,8 +194,6 @@ def set_data(self, data: TSDataSet):
         self.npt: list[int] = [f.shape[1] for f in self.flux]
         for i, time in enumerate(self.times):
             self.tms[i].set_data(time)
-        if self._nm in (NM_GP_FIXED, NM_GP_FREE):
-            self._init_gp()
 
     def _init_parameters(self) -> None:
         self.ps = ParameterSet([])
@@ -276,7 +272,7 @@ def set_gp_hyperparameters(self, sigma: float, rho: float, idata: int | None = N
         if self._gp is None:
             raise RuntimeError('The GP needs to be initialized before setting hyperparameters.')
 
-        for i in ([idata] or range(self.data.size)):
+        for i in ([idata] if idata is not None else range(self.data.size)):
             self._gp[i].kernel = terms.Matern32Term(sigma=sigma, rho=rho)
             self._gp[i].compute(self._gp_time[i], yerr=self._gp_ferr[i], quiet=True)
 

exoiris/wlpf.py

@@ -49,8 +49,8 @@ def __init__(self, tsa: TSLPF):
         fluxes, times, errors = [], [], []
         for t, f, e in zip(tsa.data.times, tsa.data.fluxes, tsa.data.errors):
             weights = where(isfinite(f) & isfinite(e), 1/e**2, 0.0)
-            mf = average(f, axis=0, weights=weights)
-            me = sqrt(1./(1./e**2).sum(0))
+            mf = average(where(isfinite(f), f, 0), axis=0, weights=weights)
+            me = sqrt(1 / weights.sum(0))
             m = isfinite(mf)
             times.append(t[m])
             fluxes.append(mf[m])
@@ -108,13 +108,13 @@ def transit_model(self, pv, copy=True):
         return self.tm.evaluate(radius_ratio, ldc, zero_epoch, period, smaxis, inclination)
 
     def optimize(self, pv0=None, method='powell', maxfev: int = 5000):
-            if pv0 is None:
-                if self.de is not None:
-                    pv0 = self.de.minimum_location
-                else:
-                    pv0 = self.ps.mean_pv
-            res = minimize(lambda pv: -self.lnposterior(pv), pv0, method=method, options={'maxfev':maxfev})
-            self._local_minimization = res
+        if pv0 is None:
+            if self.de is not None:
+                pv0 = self.de.minimum_location
+            else:
+                pv0 = self.ps.mean_pv
+        res = minimize(lambda pv: -self.lnposterior(pv), pv0, method=method, options={'maxfev':maxfev})
+        self._local_minimization = res
 
     @property
     def transit_center(self):