feat: Implement Errors-in-Variables (EIV) linear detrending (23f47885) · Commits · Gaetano Scandariato / PyAETNA

.kilo/eiv_detrending_plan.md

deleted100644 → 0

+0 −380

Changes for .kilo/eiv_detrending_plan.md: 0 added lines, 380 removed lines.
Original line number	Diff line number	Diff line
		# EIV Detrending — Kelly (2007) Hierarchical Model
		## Implementation Plan for PyAETNA

		---

		## Executive Summary & Architectural Decisions

		This document outlines the step-by-step plan for extending PyAETNA with hierarchical Errors-In-Variables (EIV) linear detrending based on Kelly (2007, ApJ 665 1489).

		### Key Architectural Decisions (Incorporating Feedback)

		1. Intercept Convention (Absorbed into Normalization):
		- Decision: No explicit EIV intercept parameter $\alpha$. The baseline offset is absorbed into the existing `normalization` parameter (stellar flux baseline for LC, systemic velocity $\gamma$ for RV).
		- Rationale: Introducing a separate intercept $\alpha$ per EIV block creates an exact linear degeneracy with `normalization` ($\gamma$), producing a flat direction in the posterior and severely degrading MCMC sampling efficiency.

		2. Latent Covariate Mixture Components ($K$ Configurable & Fully Implemented):
		- Decision: $K$ is user-configurable (`K: 1` by default). The codebase fully implements $K > 1$ using JAX-native Gaussian Mixture Models (GMM) with `Dirichlet` priors and `jax.scipy.special.logsumexp`.
		- Warning & Recommendation: $K=1$ (unimodal Gaussian prior on latent $x$) is strongly recommended. Setting $K > 1$ introduces $3K - 1$ additional hyper-parameters ($\mathbf{\pi}, \mathbf{\mu_\xi}, \mathbf{\tau}$) for the GMM, which suffers from severe MCMC label-switching, multimodal posterior degeneracies, and slower JAX compilation/sampling without providing tangible benefits for typical astronomical indicators unless the $x$ distribution is explicitly multimodal.

		3. Parameter Scope (Per-Instrument Priority):
		- Decision: Slope $\beta$ defaults to `per_instrument` scope (with `global` as an option).
		- Rationale: Detrending against ancillary indicators (e.g., $R'_{HK}$, FWHM, BIS) across different spectrographs or data types (e.g., joint LC + RV, or HARPS + ESPRESSO) must account for instrument-dependent physical sensitivities and throughput. A global slope across different instruments is physically invalid in multi-instrument contexts.

		4. Exact EIV + GP / Likelihood Integration (Phase 1 Solution):
		- Decision: Use Kelly's conditional likelihood factorization:
		$$p(x_{obs}, y \mid \dots) = p(x_{obs} \mid \mathbf{\pi}, \mathbf{\mu_\xi}, \mathbf{\tau}) \times p(y \mid x_{obs}, \dots)$$
		- Rationale: This factors the EIV model into:
		1. A marginal likelihood term for the observed covariate $x_{obs} \sim \text{GMM}(\mathbf{\pi}, \mathbf{\mu_\xi}, \mathbf{\tau}, \sigma_x^2)$.
		2. A conditional mean trend shift $\Delta \mu_{EIV}$ that enters the deterministic mean function (alongside polynomials and Fourier terms).
		3. An effective variance inflation $\sigma_{EIV, y}^2$ added in quadrature to `dy_total` (data error + jitter). When a GP is active, this effective variance seamlessly forms the diagonal noise `diag` of the GP kernel!

		5. Documentation & Configurator Directory:
		- Decision: All documentation and configurator source updates take place exclusively inside `docs/` (`public/` is autogenerated).

		---

		## Proposed Config Syntax

		The new subsection lives inside `detrending`, sibling to `polynomials`, `fourier`, and `GP`:

		```yaml
		detrending:
		eiv_linear:
		eiv_rhk:
		instruments: [HARPS] # Target instrument(s)
		x_variable: RHK # Ancillary column (RHK_ERR must also exist)
		slope: [per_instrument, [normal, 0, 5]] # Scope: per_instrument (recommended) or global
		K: 1 # Latent mixture components (default: 1)
		```

		\| Key \| Required \| Default \| Description \|
		\|-----\|----------\|---------\|-------------\|
		\| `x_variable` \| yes \| — \| Column name in data file. Paired `<NAME>_ERR` column must also exist. \|
		\| `slope` \| yes \| — \| Prior on linear slope $\beta$. Supports `per_instrument` (recommended) or `global`. \|
		\| `instruments` \| no \| all instruments \| List of instrument names to apply EIV detrending to. \|
		\| `K` \| no \| `1` \| Number of Gaussian mixture components for latent $x$. Warning issued if $K > 1$. \|

		---

		## Step-by-Step Implementation Details

		---

		### Step 1 — Math & Module Implementation (`PyAETNA/sampling/engine/eiv.py`)

		Kelly (2007) models the true latent covariate distribution as a mixture of $K$ Gaussians:
		$$\xi_i \sim \sum_{k=1}^K \pi_k \, \mathcal{N}(\mu_{\xi, k}, \tau_k^2)$$

		For $K=1$, this simplifies to a single Gaussian. For general $K \ge 1$:

		$$\gamma_{i, k} = \frac{\tau_k^2}{\tau_k^2 + \sigma_{x,i}^2}$$

		1. Marginal Log-Likelihood of $x_{obs}$:
		The component log-likelihood for data point $i$ under component $k$ is:
		$$\ln \mathcal{L}_{i, k} = \ln \pi_k - \frac{1}{2} \left[ \ln\left(2\pi (\tau_k^2 + \sigma_{x,i}^2)\right) + \frac{(x_{obs, i} - \mu_{\xi, k})^2}{\tau_k^2 + \sigma_{x,i}^2} \right]$$

		Summing across components in JAX using numerically stable `logsumexp`:
		$$\ln p(x_{obs, i}) = \text{logsumexp}_{k=1 \dots K} \left( \ln \mathcal{L}_{i, k} \right)$$

		2. Mixture Responsibilities (Weights per Point):
		$$w_{i, k} = \frac{\exp(\ln \mathcal{L}_{i, k})}{\sum_{m=1}^K \exp(\ln \mathcal{L}_{i, m})} = \text{softmax}_{k}(\ln \mathcal{L}_{i, \cdot})$$

		3. Conditional Mean Shift & Variance Inflation on $y$:
		$$\Delta \mu_{EIV, i} = \sum_{k=1}^K w_{i, k} \cdot \beta \cdot \gamma_{i, k} \cdot (x_{obs, i} - \mu_{\xi, k})$$
		$$\sigma_{EIV, y, i}^2 = \sum_{k=1}^K w_{i, k} \cdot \beta^2 \cdot \gamma_{i, k} \cdot \sigma_{x,i}^2$$

		```python
		# PyAETNA/sampling/engine/eiv.py

		import jax.numpy as jnp
		import jax.scipy.special as jsp

		def compute_eiv_components(beta, pi, mu_xi, log_tau, x_obs, x_err):
		"""
		Compute EIV conditional mean trend shift and variance inflation for y
		for K >= 1 mixture components.

		Parameters
		----------
		beta : (N,) or scalar — slope parameter
		pi : (K,) — mixture weights (sums to 1)
		mu_xi : (K,) — latent x component means
		log_tau : (K,) — log(std) of latent x components
		x_obs : (N,) — observed ancillary covariate
		x_err : (N,) — 1-sigma errors on x_obs

		Returns
		-------
		eiv_mean_shift : (N,) — trend shift to add to mean_function
		eiv_var_extra : (N,) — variance contribution to add to dy^2
		"""
		tau2 = jnp.exp(2.0 * log_tau) # (K,)
		sx2 = x_err[:, None] ** 2 # (N, 1)

		var_x = tau2[None, :] + sx2 # (N, K)
		gamma = tau2[None, :] / var_x # (N, K)

		# Component log-likelihood matrix (N, K)
		dx = x_obs[:, None] - mu_xi[None, :] # (N, K)
		log_L = jnp.log(pi)[None, :] - 0.5 * (jnp.log(2.0 * jnp.pi * var_x) + (dx ** 2) / var_x)

		# Softmax mixture responsibilities w_ik (N, K)
		w = jsp.softmax(log_L, axis=-1)

		# Weighted mean shift and variance extra (N,)
		# beta can be broadcast (N,)
		beta_col = beta if jnp.ndim(beta) > 0 else beta
		eiv_mean_shift = beta_col * jnp.sum(w * gamma * dx, axis=-1)
		eiv_var_extra = (beta_col ** 2) * jnp.sum(w * gamma * sx2, axis=-1)

		return eiv_mean_shift, eiv_var_extra


		def kelly_eiv_x_marginal_logp(pi, mu_xi, log_tau, x_obs, x_err):
		"""
		Marginal log-likelihood for observed covariate x_obs under GMM(pi, mu_xi, tau^2).
		"""
		tau2 = jnp.exp(2.0 * log_tau) # (K,)
		sx2 = x_err[:, None] ** 2 # (N, 1)
		var_x = tau2[None, :] + sx2 # (N, K)
		dx = x_obs[:, None] - mu_xi[None, :] # (N, K)

		log_L = jnp.log(pi)[None, :] - 0.5 * (jnp.log(2.0 * jnp.pi * var_x) + (dx ** 2) / var_x)
		logp_per_point = jsp.logsumexp(log_L, axis=-1) # (N,)
		return jnp.sum(logp_per_point)
		```

		---

		### Step 2 — Config Parser (`PyAETNA/sampling/config/parser/detrending.py`)

		Implement `parse_eiv_linear()`:

		```python
		def parse_eiv_linear(detrending, plModel, n_instruments, visit_ids, instrument_ids,
		free_params, fixed_params, param_scopes, param_groups,
		named_priors=None, log_func=None):
		"""Parse EIV linear detrending subsections (Kelly 2007)."""
		eiv_configs = []
		eiv_section = detrending.get('eiv_linear') or {}

		for eiv_name, eiv_cfg in eiv_section.items():
		if not isinstance(eiv_cfg, dict):
		continue

		x_variable = eiv_cfg.get('x_variable')
		if x_variable is None:
		raise ValueError(f"EIV block '{eiv_name}' is missing required 'x_variable' field")

		slope_raw = eiv_cfg.get('slope')
		if slope_raw is None:
		raise ValueError(f"EIV block '{eiv_name}' is missing required 'slope' field")

		K = int(eiv_cfg.get('K', 1))
		if K < 1:
		raise ValueError(f"EIV block '{eiv_name}': K must be >= 1, got {K}")

		if K > 1 and log_func:
		log_func(plModel, (
		f"WARNING: EIV block '{eiv_name}' sets K={K}. K > 1 introduces a "
		f"Gaussian Mixture Model for latent covariate '{x_variable}', which "
		f"can cause label-switching degeneracies and slow sampling. "
		f"K=1 is strongly recommended."
		))

		target_insts = eiv_cfg.get('instruments', plModel.instrument_names)

		# Parse slope with per_instrument or global scope
		slope_raw = resolve_config_value(
		slope_raw, named_priors, f'detrending.eiv_linear.{eiv_name}.slope'
		)
		slope_scope, slope_entries = _parse_scoped_param(
		slope_raw, ('global', 'per_instrument'),
		len(target_insts), f'{eiv_name}.slope'
		)

		base_name = f'{eiv_name}_slope'
		slope_names = _expand_scoped_param(
		slope_scope, slope_entries, base_name,
		plModel.uniqueVisits, plModel.uniqueInstruments,
		free_params, fixed_params, plModel.instrument_names,
		visit_ids=visit_ids, instrument_ids=instrument_ids
		)

		param_scopes[base_name] = slope_scope
		param_groups[base_name] = slope_names

		eiv_configs.append({
		'name': eiv_name,
		'x_variable': x_variable,
		'slope_scope': slope_scope,
		'slope_names': slope_names,
		'instruments': target_insts,
		'K': K,
		})

		return eiv_configs
		```

		---

		### Step 3 — Data Validation (`PyAETNA/sampling/engine/data_loader.py`)

		Add column check right after data loading:

		```python
		_eiv_section = config.get('detrending', {}).get('eiv_linear', {}) or {}
		for _eiv_name, _eiv_cfg in _eiv_section.items():
		if not isinstance(_eiv_cfg, dict):
		continue
		_col = _eiv_cfg.get('x_variable')
		if _col is None:
		continue
		_err_col = f'{_col}_ERR'
		if _col not in plData.dBasis:
		log_func(plModel, (
		f"ERROR: EIV block '{_eiv_name}' specifies x_variable='{_col}', "
		f"but column '{_col}' was not found in data file. "
		f"Available extra columns: {list(plData.dBasis.keys())}. Aborting."
		))
		sys.exit(1)
		if _err_col not in plData.dBasis:
		log_func(plModel, (
		f"ERROR: EIV block '{_eiv_name}' specifies x_variable='{_col}', "
		f"but required uncertainty column '{_err_col}' was not found in data file. "
		f"EIV detrending requires both '{_col}' and '{_err_col}'. Aborting."
		))
		sys.exit(1)
		log_func(plModel, f"EIV block '{_eiv_name}': validated columns '{_col}' and '{_err_col}'. OK.")
		```

		---

		### Step 4 — NumPyro Engine Integration (`PyAETNA/sampling/engine/numpyro_model.py`)

		In `_build_numpyro_model`:

		1. Pre-compute data arrays (pre-tracing):
		```python
		eiv_configs = param_config.get('eiv_configs', [])
		eiv_cache = {}
		for ecfg in eiv_configs:
		col = ecfg['x_variable']
		x_arr = np.array(plData.dBasis[col])
		xe_arr = np.array(plData.dBasis[f'{col}_ERR'])
		eiv_cache[ecfg['name']] = {
		'x_obs': jnp.array(x_arr),
		'x_err': jnp.array(xe_arr),
		'x_mean': float(np.mean(x_arr)),
		'x_std': max(float(np.std(x_arr)), 1e-9),
		}
		```

		2. Inside `model_fn` (Handling $K=1$ and $K > 1$ dynamically):
		```python
		# Evaluate EIV trend and variance inflation
		eiv_trend_total = jnp.zeros_like(pl_data.y)
		eiv_var_extra_total = jnp.zeros_like(pl_data.y)

		if eiv_configs:
		from PyAETNA.sampling.engine.eiv import compute_eiv_components, kelly_eiv_x_marginal_logp

		for ecfg in eiv_configs:
		ename = ecfg['name']
		cache = eiv_cache[ename]
		K = ecfg['K']

		if K == 1:
		pi = jnp.array([1.0])
		mu_xi = numpyro.sample(
		f'{ename}_mu_xi',
		dist.Normal(cache['x_mean'], 3.0 * cache['x_std'])
		)[None] # shape (1,)
		log_tau = numpyro.sample(
		f'{ename}_log_tau',
		dist.Uniform(np.log(0.001 * cache['x_std']), np.log(3.0 * cache['x_std']))
		)[None] # shape (1,)
		else:
		# GMM for K > 1
		pi = numpyro.sample(
		f'{ename}_pi',
		dist.Dirichlet(jnp.ones(K))
		)
		mu_xi = numpyro.sample(
		f'{ename}_mu_xi',
		dist.Normal(jnp.full(K, cache['x_mean']), 3.0 * cache['x_std'])
		)
		log_tau = numpyro.sample(
		f'{ename}_log_tau',
		dist.Uniform(jnp.full(K, np.log(0.001 * cache['x_std'])),
		jnp.full(K, np.log(3.0 * cache['x_std'])))
		)

		# Retrieve slope beta according to scope
		if ecfg['slope_scope'] == 'global':
		beta = params[ecfg['slope_names'][0]]
		else: # per_instrument
		beta = jnp.zeros_like(pl_data.y)
		for i, inst_id in enumerate(pl_model.uniqueInstruments):
		mask = pl_model.instrument_masks[i]
		beta = jnp.where(mask, params[ecfg['slope_names'][i]], beta)

		shift, var_ex = compute_eiv_components(
		beta, pi, mu_xi, log_tau, cache['x_obs'], cache['x_err']
		)
		eiv_trend_total = eiv_trend_total + shift
		eiv_var_extra_total = eiv_var_extra_total + var_ex

		# Add x-marginal log-likelihood factor
		x_logp = kelly_eiv_x_marginal_logp(pi, mu_xi, log_tau, cache['x_obs'], cache['x_err'])
		numpyro.factor(f'eiv_x_logp_{ename}', x_logp)

		# Update mean function and dy_total
		mean_function = build_mean_function(
		norm_model, planet_offset, poly_model + fourier_model + eiv_trend_total,
		plModel.normalization_mode, pl_data.instrument,
		pl_model.uniqueInstruments, data_types
		)

		# Infold EIV variance extra into dy_total
		dy_total = jnp.sqrt(dy_total ** 2 + eiv_var_extra_total)
		```

		---

		### Step 5 — Documentation (`docs/`)

		Update `docs/configuration.md` and `docs/configuration_keys.md` to detail `detrending.eiv_linear`, explaining where appropriate:
		- the math
		- Mandatory pairing of `x_variable` and `<x_variable>_ERR`.
		- `per_instrument` scope recommendation.
		- Explanation of the behavior of `mu_xi` and `log_tau` latent variables: they are optional, and if not explicitly declared, they inherit the scope of the `slope` parameter and use data-driven default priors (check `detrending.py` for more details).
		- Explanation of why intercept $\alpha$ is absorbed by normalization.
		- Full details on $K$ parameter: defaults to 1, supports $K \ge 1$ via GMM, with a clear warning explaining why $K=1$ is recommended.

		---

		### Step 6 — HTML configurator (`docs/config-generator.html`)
		Update `docs/config-generator.html` and associated files to allow the user set the EIV config section. This implementation must follow the style of the rest of the configurator.

		---

		## Verification Plan

		### Automated Tests (`tests/test_eiv.py`)
		1. `test_eiv_math_k1`: Verify `compute_eiv_components` for $K=1$ against analytical formulas.
		2. `test_eiv_math_gmm`: Verify `compute_eiv_components` and `kelly_eiv_x_marginal_logp` for $K=2, 3$ GMM using `logsumexp`.
		3. `test_eiv_parser`: Verify `parse_eiv_linear` correctly parses `K`, `per_instrument` and `global` scopes.
		4. `test_eiv_data_validation`: Test missing `<x_variable>_ERR` column triggers clean exit.
		5. `test_eiv_model_jit`: Verify JIT tracing and gradient computation of `model_fn` with EIV active for both $K=1$ and $K=2$.

		---

		## Summary of Addressed Questions

		- Q1 (Intercept): Absorbed into `normalization` to prevent severe baseline degeneracy.
		- Q2 (Mixture components K): User-configurable `K` (default 1). Fully implemented for $K \ge 1$ using JAX-native Dirichlet-GMM and `logsumexp`, accompanied by a prominent warning that $K=1$ is recommended.
		- Q3 (Slope Scope): `per_instrument` implemented as primary scope (with `global` supported).
		- Q4 (EIV + GP interaction): Solved via exact conditional decomposition (EIV trend in mean function, EIV variance in `dy_total`/GP diagonal).
		- Q5 (Configurator directory): All documentation edits targeted strictly to `docs/`.

.kilo/eiv_status.md

deleted100644 → 0

+0 −28

Changes for .kilo/eiv_status.md: 0 added lines, 28 removed lines.
Original line number	Diff line number	Diff line
		# EIV Detrending Integration Status

		## Premise

		In the rpoerting phase of pyaetna, if a polynomial detrending is activated then the bestfit model takes the polynomial model into account, removes it from the observed data and is shown in the bestfit plot model as a non time trend if the x_variable is not BJD. in the end, the difference between EIV and linear polynomial detrending is that, during MC sampling, the slope of the trend is estimated in a mor robust statistical way than standard linear fit. after that, the detrending in the best fit model (and thus the whole reporting phase) should behave the same

		## Objective
		Ensure EIV (Errors-In-Variables) detrending is reported and plotted consistently with linear polynomial detrending:
		- Accounted for in the best-fit model.
		- Removed from observed data (residuals).
		- Identified as a non-time trend if the x-variable is not BJD.

		## Work Completed
		1. Model Components: Updated `PyAETNA/reporting/model_components.py` to reconstruct EIV mean shifts and export both `eiv_terms` and `eiv_configs` in the return dictionary.
		2. Math Helpers: Updated `PyAETNA/reporting/model_math.py`'s `_compute_nontime_trend` to accept and process EIV terms based on their configuration.
		3. Data Preparation: Updated `PyAETNA/reporting/plot_data_prep.py` to use `model_comp['eiv_configs']` when calculating non-time trends for both the full data and the fine-grid aliases.
		4. Confidence Bands: Fixed `TypeError` and `KeyError` in `PyAETNA/reporting/confidence_bands.py` to ensure EIV components are passed correctly during MCMC sample processing.

		## Open Issues
		Despite the implementation of the above, the following symptoms persist:
		- Missing Plots: In the bestfit panles the nontime detrend is not showing the EIV contribution.
		- Broad Residuals: The residuals remain too broad, suggesting that the EIV trend is not being successfully subtracted from the observed data before plotting.

		## Next Steps for Investigation
		- Verify the reporting flow to ensure that the nontime detrend in the EIV way is treated as the polynomial one

		## Bonus track
		Once the above is fixed, one minor thing is left. When `_ERR` columns are detected in the `dBasis` dictionnary, the code is already triggered to plot the correlation plots as errorbars, not scatter plots. Similarly, when the trigger is activated, the linear fit shown as a red solid line should be computed using the slope estimated by the `kelly_2007_quick_estimate` function inside `/dati/ricerca/gaps/TOI5572/readData.py`

PyAETNA.egg-info/PKG-INFO

+1 −0

Changes for PyAETNA.egg-info/PKG-INFO: 1 added line, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -301,6 +301,7 @@ This project is licensed under the MIT License — see the [LICENSE](LICENSE) fi
		- check the checkpoint_interval logic
		- implement different limb darkening laws
		- better arrangement of legends and text in plots
		- remove double INFO messages in the log

		## Contact

PyAETNA.egg-info/SOURCES.txt

+3 −0

Changes for PyAETNA.egg-info/SOURCES.txt: 3 added lines, 0 removed lines.
Original line number	Diff line number	Diff line
		@@ -8,6 +8,7 @@ PyAETNA/gp.py
		PyAETNA/model_utils.py
		PyAETNA/planet_phys.py
		PyAETNA/run.py
		PyAETNA/version.py
		PyAETNA.egg-info/PKG-INFO
		PyAETNA.egg-info/SOURCES.txt
		PyAETNA.egg-info/dependency_links.txt
		@@ -62,6 +63,7 @@ PyAETNA/sampling/config/parser/scoping.py
		PyAETNA/sampling/config/parser/star.py
		PyAETNA/sampling/engine/checkpoint.py
		PyAETNA/sampling/engine/data_loader.py
		PyAETNA/sampling/engine/eiv.py
		PyAETNA/sampling/engine/logging.py
		PyAETNA/sampling/engine/numpyro_model.py
		PyAETNA/sampling/engine/runner.py
		@@ -73,6 +75,7 @@ PyAETNA/sampling/model/gp_models.py
		PyAETNA/scripts/normalize_flux_by_visit.py
		PyAETNA/scripts/plot_random_gp.py
		tests/test_behavioral_ties.py
		tests/test_eiv.py
		tests/test_exptime_integration.py
		tests/test_gp_per_visit.py
		tests/test_prior_refs.py

PyAETNA/reporting/confidence_bands.py

+6 −2

Changes for PyAETNA/reporting/confidence_bands.py: 6 added lines, 2 removed lines.
Original line number	Diff line number	Diff line
		@@ -527,7 +527,8 @@ def compute_grid_bands(
		if has_non_time_trends:
		_ov_tr_nontime_grid = _compute_nontime_trend(
		model_comp_grid['poly_terms'], model_comp_grid['fourier_terms'],
		poly_cfgs, four_cfgs, n_grid)
		model_comp_grid['eiv_terms'],
		poly_cfgs, four_cfgs, model_comp_grid['eiv_configs'], n_grid)
		_ov_tr_time_grid = tr_grid - _ov_tr_nontime_grid
		if data_type_in_mask == 'RV' or normalization_mode == 'additive':
		_ov_total_grid = total_grid - _ov_tr_nontime_grid
		@@ -643,7 +644,10 @@ class ConfidenceBandComputer:
		param_config, skip_gp=True)
		result['nontime_trend'] = _compute_nontime_trend(
		comp_nt['poly_terms'], comp_nt['fourier_terms'],
		poly_cfgs, four_cfgs, len(plData.time)
		comp_nt['eiv_terms'],
		poly_cfgs, four_cfgs,
		comp_nt['eiv_configs'],
		len(plData.time)
		)

		return result