refactor of group 2 (gp): improve GP numerical stability and JIT compatibility

PyAETNA GP Module - Shared Gaussian Process utilities

This module provides GP kernel building and prediction functions used by both

the sampler (PyAETNA/sampling/engine/runner.py) and plotting (PyAETNA/reporting/lightcurve_plots.py) modules.

the sampler (PyAETNA/sampling/engine/runner.py) and reporting modules.

Centralizing GP code here ensures consistency between sampling and visualization.

See ``docs/gp_per_visit_analysis/GP_PER_VISIT_ANALYSIS.md`` for benchmarks.

"""

from functools import reduce

import jax

import jax.numpy as jnp

import numpy as np

from tinygp import GaussianProcess, kernels, transforms

from tinygp import GaussianProcess

# Import kernel functions from the new dedicated kernels/factory.py module

from PyAETNA.kernels.factory import (

# GP Prediction (JIT-compiled for speed)

# =============================================================================

@jax.jit

def gp_predict(kernel, x, y, dy_sq, x_test, log_func=None, log_target=None):

    """

    JIT-compiled GP prediction.

    GP prediction wrapper. Handles shape validation and logging before 

    calling the JIT-compiled core.

    Args:

        kernel: tinygp kernel object

    Returns:

        mean, std: Predicted mean and standard deviation at x_test

    """

    # Convert arrays to ensure proper types

    # Convert to arrays for shape checking

    x = jnp.asarray(x)

    y = jnp.asarray(y)

    dy_sq = jnp.asarray(dy_sq)

    x_test = jnp.asarray(x_test)

    # Ensure dy_sq is 1D and has the correct length - this solves a specific issue in 

    # some kernel implementations where dy_sq could be inadvertently passed as matrix [n,n]  

    # instead of vector [n]. For tinygp compatibility, it must be a 1D array.

    # Shape repair and validation

    if dy_sq.ndim > 1:

        try:

            # Try to flatten the array; this is the typical fix 

            flat = dy_sq.flatten()

            if len(flat) == len(y):

        flat = dy_sq.reshape(-1)

        if flat.shape[0] == y.shape[0]:

            dy_sq = flat

        else:

                # If that doesn't work or shape mismatches, make it right (default behavior)

                # This should not happen - only for very weird edge cases in the test workflow

            msg = f"ERROR: dy_sq shape {dy_sq.shape} cannot be flattened to match y length {y.shape[0]}"

            if log_func:

                    log_func(log_target, f"WARNING: dy_sq has wrong shape {dy_sq.shape} != {len(y)}, using defaults")

                dy_sq = jnp.full(len(y), 1e-6)  

        except Exception as e:

            # On any error during shape conversion, use default safe fallback

            if log_func:

                log_func(log_target, f"WARNING: Error in shape conversion of dy_sq: {e}, using defaults")

            dy_sq = jnp.full(len(y), 1e-6)

                log_func(log_target, msg)

            raise ValueError(msg)

    # Final sanity check 

    if dy_sq.ndim != 1 or len(dy_sq) != len(y):

    if dy_sq.shape[0] != y.shape[0]:

        msg = f"ERROR: dy_sq length {dy_sq.shape[0]} must match y length {y.shape[0]}"

        if log_func:

            log_func(log_target, f"WARNING: dy_sq has wrong final shape {dy_sq.shape} != {len(y)}, using defaults")

        dy_sq = jnp.full(len(y), 1e-6)

            log_func(log_target, msg)

        raise ValueError(msg)

    # Extra debugging for tracing issues

    # Before creating the GaussianProcess, ensure we're in safe numerical state

    if len(dy_sq) != len(y):

        if log_func:

            log_func(log_target, f"ERROR: dy_sq ({len(dy_sq)}) and y ({len(y)}) lengths inconsistent")

        dy_sq = jnp.full(len(y), 1e-6)

    return _gp_predict_jit(kernel, x, y, dy_sq, x_test)

@jax.jit

def _gp_predict_jit(kernel, x, y, dy_sq, x_test):

    """Core JIT-compiled GP prediction logic."""

    gp = GaussianProcess(kernel, x, diag=dy_sq)

    gp_cond = gp.condition(y, x_test)

    return gp_cond.gp.loc, jnp.sqrt(gp_cond.gp.variance)

"""

PyAETNA kernel factory - builds tinygp kernel objects from configuration.

This module contains the low-level (`build_single_kernel`) and high-level

(`build_gp_kernel_from_specs` and its `build_gp_kernel_for_plotting` /

`build_gp_kernel_for_sampling` wrappers) kernel construction functions used

throughout PyAETNA's GP sampling and reporting code (see `PyAETNA/gp.py`).

"""

import jax.numpy as jnp

import numpy as np

from tinygp import kernels, transforms

import equinox as eqx

from functools import reduce

from typing import Any

import re

SUPPORTED_KERNELS = (

                        # (Don't call tolist() as it causes JAX concretization errors during tracing)

                        param_values.append(val)

                    # Check for the possibility of multi-valued parameter expansion (when we might have list inputs from YAML)  

                    if full_pname in params and isinstance(params[full_pname], list):

                        # This is a special case - when YAML defines parameters as lists like 

                        # `log_gp_amp: [ ['uniform', -3, 1], ['uniform', -4, 1], ['uniform', -3, 1] ]`

                        # These are not actually array parameters but sets of parameters

                        # and should be processed directly by NumPyro as separate parameters, so no action needed here.

                        pass

        # Build kernel term

        kernel_term = build_single_kernel(kernel_type, param_values, subspace_idx, use_jax=use_jax)

        kernel_terms.append(kernel_term)

    _rename_params_with_instrument_names,

)

from .plot_models import plot_models

from .model_components import compute_model_components, write_model_components

from .model_components import compute_model_components, save_model_components_binary

from .model_math import (

    _compute_planet_model_for_plotting,

    _apply_scoped_values,

# -*- coding: utf-8 -*-

"""

lightcurve_plots.py  –  COMPATIBILITY SHIM

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

This module used to contain all plotting helpers.  It has been split into

focused sub-modules.  All public names are re-exported here for backwards

compatibility with external code that imports from ``lightcurve_plots``.

This shim will be removed in a future clean-up commit.

"""

from .plot_models import plot_models

from .model_components import compute_model_components, write_model_components

from .model_math import (

    _compute_planet_model_for_plotting,

    _apply_scoped_values,

    _precompute_sample_data,

    _compute_nontime_trend,

)

from .noise_plots import plot_noise_budget

from .periodogram_plots import (

    plot_residuals_histogram,

    plot_periodograms_comparison,

    plot_periodogram,

    plot_correlated_noise,

    plot_residuals_vs_dbasis,

)

from .phase_plots import (

    _time_to_phase,

    plot_phase_folded,

    plot_phase_folded_rv,

    plot_phase_folded_rv_rm_zoom,

)

from .ttv_plots import (

    plot_ttv_oc,

    plot_ttv_per_transit,

    plot_ttv_transit_overlay,

)

__all__ = [

    'plot_models',

    'compute_model_components', 'write_model_components',

    '_compute_planet_model_for_plotting', '_apply_scoped_values',

    '_precompute_sample_data', '_compute_nontime_trend',

    'plot_noise_budget',

    'plot_residuals_histogram', 'plot_periodograms_comparison',

    'plot_periodogram', 'plot_correlated_noise', 'plot_residuals_vs_dbasis',

    '_time_to_phase', 'plot_phase_folded',

    'plot_phase_folded_rv', 'plot_phase_folded_rv_rm_zoom',

    'plot_ttv_oc', 'plot_ttv_per_transit', 'plot_ttv_transit_overlay',

]

Functions

---------

compute_model_components

write_model_components

save_model_components_binary

"""

import numpy as np

import jax

            gp_mean[inst_mask] = mu_full_np

            gp_std[inst_mask]  = std_full_np

            # Record per-kernel-term labels (needed by write_model_components).

            # Record per-kernel-term labels (needed by save_model_components_binary).

            # When there is only one kernel term the composite output IS the

            # single-term output; for multi-term configs we label the

            # composite as the first (dominant) term — per-term decomposition

    }

def write_model_components(outfile, components, plData=None, mask=None, sort_by_time=True):

    """Write model component columns to a whitespace-separated file.

    The file columns are in this order:

        [planet terms...], [polynomial terms...], [fourier terms...], [gp terms...], RESIDUALS

    The observed data columns (BJD, Y, SIGMA) are NOT repeated here;

    they are already present in the original data file.

    Args:

        outfile: Path to output file

        components: Dict returned by `compute_model_components`

        plData: Optional PlData object (used to compute BJD if not present in components)

        mask: Optional boolean mask or index array to subset rows (e.g., per-visit)

        sort_by_time: Whether to sort rows by BJD (default True)

    """

    # Gather core arrays

    bjd = components.get('bjd', None)

    if bjd is None and plData is not None and hasattr(plData, 'time'):

        bjd = np.array(plData.time) + getattr(plData, 'timeCenter', 0.0)

    y = components.get('y', None)

    dy = components.get('dy', None)

    total = components.get('total', None)

    residuals = components.get('residuals', None)

    n_pts = len(bjd) if bjd is not None else len(y)

    if mask is None:

        sel = np.arange(n_pts)

    else:

        sel = np.where(mask)[0] if isinstance(mask, (np.ndarray, list)) else np.arange(n_pts)[mask]

    # Collect term dicts

    poly_terms = components.get('poly_terms', {})

    fourier_terms = components.get('fourier_terms', {})

    gp_terms = components.get('gp_terms', {})

    # Determine column order (place planet columns first)

    planet_terms = components.get('planet_terms', {})

    planet_keys = list(planet_terms.keys())

    poly_keys = sorted(list(poly_terms.keys()))

    fourier_keys = sorted(list(fourier_terms.keys()))

    gp_keys = sorted(list(gp_terms.keys()))

    # Do not include BJD/Y/SIGMA in the output file (they are present in original data)

    header = planet_keys + poly_keys + fourier_keys + gp_keys + ['RESIDUALS']

    # Build columns (planet columns first, then components, then residuals)

    cols = []

    for k in planet_keys:

        cols.append(np.array(planet_terms[k])[sel])

    for k in poly_keys:

        cols.append(np.array(poly_terms[k])[sel])

    for k in fourier_keys:

        cols.append(np.array(fourier_terms[k])[sel])

    for k in gp_keys:

        cols.append(np.array(gp_terms[k])[sel])

    # Residuals

    cols.append(np.array(residuals)[sel])

    # Stack

    data = np.column_stack(cols)

    if sort_by_time:

        if bjd is None:

            _log('WARNING: BJD not available; skipping time-based sorting for components file')

        else:

            sort_idx = np.argsort(np.array(bjd)[sel])

            data = data[sort_idx, :]

    # Format output with equal-width columns for readability

    # Choose column width based on longest header name and numeric width

    n_cols = len(header)

    longest_name = max(len(h) for h in header)

    value_width = 13  # default value field width

    col_width = max(longest_name, value_width) + 2  # padding

    # Write file with fixed-width columns

    try:

        with open(outfile, 'w') as f:

            # Header as comment line with fixed-width columns

            f.write('# ' + ''.join([h.ljust(col_width) for h in header]) + '\n')

            fmt = f'{{:>{col_width}.6g}}'  # general format, right-aligned into width

            for row in data:

                f.write(''.join([fmt.format(v) for v in row]) + '\n')

        _log(f"Wrote model component file: {outfile}")

    except Exception as e:

        _log(f"ERROR: Failed to write model component file {outfile}: {str(e)}")

def save_model_components_binary(outfile, map_data, median_data, quantile_data, plData, plModel, param_config):

    """

    Save model components (MAP, Median, Quantiles) to a binary pickle file.