Commit ae6d6fec authored by Gaetano Scandariato's avatar Gaetano Scandariato
Browse files

Refactor of group 4 (reporting and visualization) completed. 

- Remove unused imports from reporting modules

- Replace bare except: clauses with except Exception as e and add warning logs

- Convert silent YAML reload failure in pdf_report.py into logged fallback behavior

- Replace stray print() logging in ttv_plots.py with module logger

- Optimize confidence band computation by reducing redundant np.percentile() calls

- Smoke-test reporting changes via generate_report() and direct periodogram plotting

- Introduce ModelComponentsConfig dataclass to simplify compute_model_components signature and update all call sites.

- Optimize compute_grid_bands by vectorizing np.percentile calls to reduce sorting overhead.

- Establish testing foundation for the reporting package, including synthetic fixtures, unit tests for core computation functions, and a smoke test for the PDF report generation pipeline.
parent 7222828b

PyAETNA/reporting/confidence_bands.py

+97 −57
Changes for PyAETNA/reporting/confidence_bands.py: 97 added lines, 57 removed lines.
Original line number Diff line number Diff line
@@ -14,7 +14,7 @@ import numpy as np 
import time as _timer



from PyAETNA.model_utils import build_partial_for_plot, build_composite_for_plot

from .model_components import compute_model_components

from .model_components import compute_model_components, ModelComponentsConfig

from .model_math import _compute_nontime_trend

from .gp_plotting import _precompute_gp_grid_context, _gp_predict_on_grid
@@ -103,8 +103,14 @@ def compute_data_bands( 
            c_total = sample_cache['total'][ii]

            if has_non_time_trends:

                _comp_nt = compute_model_components(

                    sample_cache['params'][ii], plData, plModel,

                    param_config, skip_gp=True)

                    ModelComponentsConfig(

                        params=sample_cache['params'][ii],

                        plData=plData,

                        plModel=plModel,

                        param_config=param_config,

                        skip_gp=True

                    )

                )

                nontime_trend_data_samples.append(

                    _compute_nontime_trend(

                        _comp_nt['poly_terms'], _comp_nt['fourier_terms'],
@@ -114,8 +120,15 @@ def compute_data_bands( 
            params_i = {name: samples_dict[name][sample_indices[ii]]

                        for name in samples_dict.keys()}

            params_i.update(param_config['fixed_params'])

            comp_i = compute_model_components(params_i, plData, plModel, param_config,

                                              skip_gp=True)

            comp_i = compute_model_components(

                ModelComponentsConfig(

                    params=params_i,

                    plData=plData,

                    plModel=plModel,

                    param_config=param_config,

                    skip_gp=True

                )

            )

            cn, cp, ct, c_total = (comp_i['norm'], comp_i['planet'],

                                   comp_i['trend'], comp_i['total'])

            if has_non_time_trends:
@@ -227,8 +240,6 @@ def compute_grid_bands( 
    visit_to_plot,

    log_func,

):

    if log_func is None:

        log_func = lambda msg: None

    """Compute smooth model bands on the plotting time grid.



    Uses the law of total variance to cleanly separate parameter uncertainty
@@ -277,6 +288,8 @@ def compute_grid_bands( 
        ``_ov_pt_q{16,84}_grid``, ``_ov_tot_q{16,84}_grid``,

        ``_ov_mt_q{16,84}_grid``.

    """

    if log_func is None:

        log_func = lambda msg: None

    # Initialise all band variables to None (populated in the branches below)

    pn_q16_grid = pn_q84_grid = None

    pt_q16_grid = pt_q84_grid = None
@@ -337,8 +350,15 @@ def compute_grid_bands( 
                            for name in samples_dict.keys()}

                params_i.update(param_config['fixed_params'])



            comp_i = compute_model_components(params_i, plData_grid, plModel, param_config,

                                              skip_gp=True)

            comp_i = compute_model_components(

                ModelComponentsConfig(

                    params=params_i,

                    plData=plData_grid,

                    plModel=plModel,

                    param_config=param_config,

                    skip_gp=True

                )

            )

            norm_samples_grid.append(comp_i['norm'])

            pm_samples_grid.append(comp_i['planet'])

            tr_samples_grid.append(comp_i['trend'])
@@ -354,7 +374,14 @@ def compute_grid_bands( 
                    dy_i  = sample_cache['dy_total'][ii]

                else:

                    comp_real_i = compute_model_components(

                        params_i, plData, plModel, param_config, skip_gp=True)

                        ModelComponentsConfig(

                            params=params_i,

                            plData=plData,

                            plModel=plModel,

                            param_config=param_config,

                            skip_gp=True

                        )

                    )

                    res_i = np.array(comp_real_i['y']) - np.array(comp_real_i['mean_function'])

                    dy_i  = np.array(comp_real_i['dy_total'])
@@ -394,12 +421,12 @@ def compute_grid_bands( 
            planet_norm_samples_grid  = norm_samples_grid + pm_samples_grid

            planet_trend_samples_grid = norm_samples_grid + pm_samples_grid + tr_samples_grid

            planet_trend_gp_loc_grid  = planet_trend_samples_grid + gp_loc_samples_grid

            pn_q16_grid  = np.percentile(planet_norm_samples_grid,  16, axis=0)

            pn_q84_grid  = np.percentile(planet_norm_samples_grid,  84, axis=0)

            pt_q16_grid  = np.percentile(planet_trend_samples_grid, 16, axis=0)

            pt_q84_grid  = np.percentile(planet_trend_samples_grid, 84, axis=0)

            tot_q16_grid = np.percentile(planet_trend_gp_loc_grid,  16, axis=0) - gp_proc_std

            tot_q84_grid = np.percentile(planet_trend_gp_loc_grid,  84, axis=0) + gp_proc_std

            pn_q16, pn_q84 = np.percentile(planet_norm_samples_grid, [16, 84], axis=0)

            pn_q16_grid, pn_q84_grid = pn_q16, pn_q84

            pt_q16, pt_q84 = np.percentile(planet_trend_samples_grid, [16, 84], axis=0)

            pt_q16_grid, pt_q84_grid = pt_q16, pt_q84

            tot_q16, tot_q84 = np.percentile(planet_trend_gp_loc_grid, [16, 84], axis=0)

            tot_q16_grid, tot_q84_grid = tot_q16 - gp_proc_std, tot_q84 + gp_proc_std



        elif not is_rv_grid and has_planets:

            _z = np.zeros_like(pm_samples_grid)
@@ -410,12 +437,12 @@ def compute_grid_bands( 
            planet_trend_gp_loc_grid  = build_composite_for_plot(

                norm_samples_grid, pm_samples_grid, tr_samples_grid,

                gp_loc_samples_grid, normalization_mode, 'LC')

            pn_q16_grid  = np.percentile(planet_norm_samples_grid,  16, axis=0)

            pn_q84_grid  = np.percentile(planet_norm_samples_grid,  84, axis=0)

            pt_q16_grid  = np.percentile(planet_trend_samples_grid, 16, axis=0)

            pt_q84_grid  = np.percentile(planet_trend_samples_grid, 84, axis=0)

            tot_q16_grid = np.percentile(planet_trend_gp_loc_grid,  16, axis=0) - gp_proc_std

            tot_q84_grid = np.percentile(planet_trend_gp_loc_grid,  84, axis=0) + gp_proc_std

            pn_q16, pn_q84 = np.percentile(planet_norm_samples_grid, [16, 84], axis=0)

            pn_q16_grid, pn_q84_grid = pn_q16, pn_q84

            pt_q16, pt_q84 = np.percentile(planet_trend_samples_grid, [16, 84], axis=0)

            pt_q16_grid, pt_q84_grid = pt_q16, pt_q84

            tot_q16, tot_q84 = np.percentile(planet_trend_gp_loc_grid, [16, 84], axis=0)

            tot_q16_grid, tot_q84_grid = tot_q16 - gp_proc_std, tot_q84 + gp_proc_std



        else:

            if is_rv_grid:
@@ -428,22 +455,22 @@ def compute_grid_bands( 
                mean_trend_gp_loc_grid  = build_composite_for_plot(

                    norm_samples_grid, _z, tr_samples_grid,

                    gp_loc_samples_grid, normalization_mode, 'LC')

            mt_q16_grid  = np.percentile(mean_trend_samples_grid, 16, axis=0)

            mt_q84_grid  = np.percentile(mean_trend_samples_grid, 84, axis=0)

            tot_q16_grid = np.percentile(mean_trend_gp_loc_grid,  16, axis=0) - gp_proc_std

            tot_q84_grid = np.percentile(mean_trend_gp_loc_grid,  84, axis=0) + gp_proc_std

            mt_q16, mt_q84 = np.percentile(mean_trend_samples_grid, [16, 84], axis=0)

            mt_q16_grid, mt_q84_grid = mt_q16, mt_q84

            tot_q16, tot_q84 = np.percentile(mean_trend_gp_loc_grid, [16, 84], axis=0)

            tot_q16_grid, tot_q84_grid = tot_q16 - gp_proc_std, tot_q84 + gp_proc_std



        if has_planets:

            pm_q16_grid = np.percentile(pm_samples_grid, 16, axis=0)

            pm_q84_grid = np.percentile(pm_samples_grid, 84, axis=0)

            pm_q16, pm_q84 = np.percentile(pm_samples_grid, [16, 84], axis=0)

            pm_q16_grid, pm_q84_grid = pm_q16, pm_q84



        norm_q16_grid = np.percentile(norm_samples_grid, 16, axis=0)

        norm_q84_grid = np.percentile(norm_samples_grid, 84, axis=0)

        norm_q16, norm_q84 = np.percentile(norm_samples_grid, [16, 84], axis=0)

        norm_q16_grid, norm_q84_grid = norm_q16, norm_q84



        # GP bands: centered around 0 (pure GP uncertainty)

        _gp_only_samples  = gp_loc_samples_grid - gp_grid

        gp_centred_q16_grid  = np.percentile(_gp_only_samples, 16, axis=0) - gp_proc_std

        gp_centred_q84_grid  = np.percentile(_gp_only_samples, 84, axis=0) + gp_proc_std

        gp_q16, gp_q84 = np.percentile(_gp_only_samples, [16, 84], axis=0)

        gp_centred_q16_grid, gp_centred_q84_grid = gp_q16 - gp_proc_std, gp_q84 + gp_proc_std



        if has_non_time_trends:

            # Compute per-component bands on the grid
@@ -456,9 +483,8 @@ def compute_grid_bands( 
            

            for name in comp_names:

                comp_samples = np.array([s[name] for s in tr_nontime_samples_grid])

                tr_nt_q16_grid[name] = np.percentile(comp_samples, 16, axis=0)

                tr_nt_q84_grid[name] = np.percentile(comp_samples, 84, axis=0)

                tr_nt_median_grid[name] = np.percentile(comp_samples, 50, axis=0)

                q16, q50, q84 = np.percentile(comp_samples, [16, 50, 84], axis=0)

                tr_nt_q16_grid[name], tr_nt_median_grid[name], tr_nt_q84_grid[name] = q16, q50, q84

            

            # Also compute the aggregate non-time trend for overlay calculations

            # tr_nontime_total_samples_grid: (n_samples, n_grid)
@@ -469,15 +495,15 @@ def compute_grid_bands( 
            _tr_time_centred_samples = build_partial_for_plot(

                norm_samples_grid, _z_band, tr_samples_grid,

                normalization_mode, 'RV' if is_rv_grid else 'LC')

            tr_time_centred_q16_grid = np.percentile(_tr_time_centred_samples, 16, axis=0)

            tr_time_centred_q84_grid = np.percentile(_tr_time_centred_samples, 84, axis=0)

            tr_q16, tr_q84 = np.percentile(_tr_time_centred_samples, [16, 84], axis=0)

            tr_time_centred_q16_grid, tr_time_centred_q84_grid = tr_q16, tr_q84



        if is_rv_grid:

            _detr_mf_samples = pm_samples_grid

        else:

            _detr_mf_samples = 1.0 + pm_samples_grid

        mf_q16_grid = np.percentile(_detr_mf_samples, 16, axis=0)

        mf_q84_grid = np.percentile(_detr_mf_samples, 84, axis=0)

        mf_q16, mf_q84 = np.percentile(_detr_mf_samples, [16, 84], axis=0)

        mf_q16_grid, mf_q84_grid = mf_q16, mf_q84



        if has_non_time_trends:

            _ov_tr_time_samples = tr_samples_grid - tr_nontime_total_samples_grid
@@ -493,16 +519,16 @@ def compute_grid_bands( 
            _tr_time_centred_samples = build_partial_for_plot(

                norm_samples_grid, _z_band, _ov_tr_time_samples,

                normalization_mode, 'RV' if is_rv_grid else 'LC')

            tr_time_centred_q16_grid = np.percentile(_tr_time_centred_samples, 16, axis=0)

            tr_time_centred_q84_grid = np.percentile(_tr_time_centred_samples, 84, axis=0)

            tr_q16, tr_q84 = np.percentile(_tr_time_centred_samples, [16, 84], axis=0)

            tr_time_centred_q16_grid, tr_time_centred_q84_grid = tr_q16, tr_q84



            if is_rv_grid and has_planets:

                _ov_pt_samples    = norm_samples_grid + pm_samples_grid + _ov_tr_time_samples

                _ov_pt_gp         = _ov_pt_samples + gp_loc_samples_grid

                _ov_pt_q16_grid   = np.percentile(_ov_pt_samples, 16, axis=0)

                _ov_pt_q84_grid   = np.percentile(_ov_pt_samples, 84, axis=0)

                _ov_tot_q16_grid  = np.percentile(_ov_pt_gp, 16, axis=0) - gp_proc_std

                _ov_tot_q84_grid  = np.percentile(_ov_pt_gp, 84, axis=0) + gp_proc_std

                ov_pt_q16, ov_pt_q84 = np.percentile(_ov_pt_samples, [16, 84], axis=0)

                _ov_pt_q16_grid, _ov_pt_q84_grid = ov_pt_q16, ov_pt_q84

                ov_tot_q16, ov_tot_q84 = np.percentile(_ov_pt_gp, [16, 84], axis=0)

                _ov_tot_q16_grid, _ov_tot_q84_grid = ov_tot_q16 - gp_proc_std, ov_tot_q84 + gp_proc_std

            elif not is_rv_grid and has_planets:

                _ov_pt_samples    = build_partial_for_plot(

                    norm_samples_grid, pm_samples_grid, _ov_tr_time_samples,
@@ -510,10 +536,10 @@ def compute_grid_bands( 
                _ov_pt_gp         = build_composite_for_plot(

                    norm_samples_grid, pm_samples_grid, _ov_tr_time_samples,

                    gp_loc_samples_grid, normalization_mode, 'LC')

                _ov_pt_q16_grid   = np.percentile(_ov_pt_samples, 16, axis=0)

                _ov_pt_q84_grid   = np.percentile(_ov_pt_samples, 84, axis=0)

                _ov_tot_q16_grid  = np.percentile(_ov_pt_gp, 16, axis=0) - gp_proc_std

                _ov_tot_q84_grid  = np.percentile(_ov_pt_gp, 84, axis=0) + gp_proc_std

                ov_pt_q16, ov_pt_q84 = np.percentile(_ov_pt_samples, [16, 84], axis=0)

                _ov_pt_q16_grid, _ov_pt_q84_grid = ov_pt_q16, ov_pt_q84

                ov_tot_q16, ov_tot_q84 = np.percentile(_ov_pt_gp, [16, 84], axis=0)

                _ov_tot_q16_grid, _ov_tot_q84_grid = ov_tot_q16 - gp_proc_std, ov_tot_q84 + gp_proc_std

            else:

                if is_rv_grid:

                    _ov_mt_samples = norm_samples_grid + _ov_tr_time_samples
@@ -526,10 +552,10 @@ def compute_grid_bands( 
                    _ov_mt_gp = build_composite_for_plot(

                        norm_samples_grid, _z, _ov_tr_time_samples,

                        gp_loc_samples_grid, normalization_mode, 'LC')

                _ov_mt_q16_grid  = np.percentile(_ov_mt_samples, 16, axis=0)

                _ov_mt_q84_grid  = np.percentile(_ov_mt_samples, 84, axis=0)

                _ov_tot_q16_grid = np.percentile(_ov_mt_gp, 16, axis=0) - gp_proc_std

                _ov_tot_q84_grid = np.percentile(_ov_mt_gp, 84, axis=0) + gp_proc_std

                ov_mt_q16, ov_mt_q84 = np.percentile(_ov_mt_samples, [16, 84], axis=0)

                _ov_mt_q16_grid, _ov_mt_q84_grid = ov_mt_q16, ov_mt_q84

                ov_tot_q16, ov_tot_q84 = np.percentile(_ov_mt_gp, [16, 84], axis=0)

                _ov_tot_q16_grid, _ov_tot_q84_grid = ov_tot_q16 - gp_proc_std, ov_tot_q84 + gp_proc_std



    else:

        # Fallback: no samples — use best_params only (no bands)
@@ -636,8 +662,15 @@ class ConfidenceBandComputer: 
                name: self.samples_dict[name][sample_indices[sample_idx]]

                for name in self.samples_dict.keys()

            }

            components = compute_model_components(params_i, plData, plModel,

                                                 param_config, skip_gp=True)

            components = compute_model_components(

                ModelComponentsConfig(

                    params=params_i,

                    plData=plData,

                    plModel=plModel,

                    param_config=param_config,

                    skip_gp=True

                )

            )

        

        # Build per-point composites (RV vs LC)

        cp = components['planet']
@@ -662,8 +695,15 @@ class ConfidenceBandComputer: 
        

        # Optionally compute non-time trend

        if has_non_time_trends:

            comp_nt = compute_model_components(params_i, plData, plModel,

                                              param_config, skip_gp=True)

            comp_nt = compute_model_components(

                ModelComponentsConfig(

                    params=params_i,

                    plData=plData,

                    plModel=plModel,

                    param_config=param_config,

                    skip_gp=True

                )

            )

            result['nontime_trend'] = _compute_nontime_trend(

                comp_nt['poly_terms'], comp_nt['fourier_terms'],

                comp_nt['eiv_terms'],

PyAETNA/reporting/model_components.py

+46 −11
Changes for PyAETNA/reporting/model_components.py: 46 added lines, 11 removed lines.
Original line number Diff line number Diff line
@@ -31,22 +31,38 @@ from .model_math import ( 
)



from PyAETNA.sampling.engine.eiv import compute_eiv_components

from dataclasses import dataclass





def compute_model_components(params, plData, plModel, param_config, skip_gp=False,

                             skip_rm=False):

@dataclass

class ModelComponentsConfig:

    """

    Configuration object for compute_model_components to avoid long parameter lists.

    """

    params: dict

    plData: Any

    plModel: Any

    param_config: dict

    skip_gp: bool = False

    skip_rm: bool = False



def compute_model_components(

    config_or_params: Any,

    plData: Any = None,

    plModel: Any = None,

    param_config: Any = None,

    skip_gp: bool = False,

    skip_rm: bool = False,

):

    """

    Compute individual model components for plotting (scope-aware).

    

    Args:

        params: Dict of all parameter values (free + fixed)

        plData: PlData object

        plModel: PlModel object

        param_config: Parameter configuration (v3.0 format with scope metadata)

        skip_gp: If True, skip all GP computation and return zeros for GP arrays.

            Use this in sample-band loops where the GP output is superseded by a

            separate grid-based GP computation (avoids expensive per-visit tinygp

            solves on the full data array for each posterior sample).

        config_or_params: ModelComponentsConfig object OR params dictionary.

        plData: PlData object (when config_or_params is params dict).

        plModel: PlModel object (when config_or_params is params dict).

        param_config: param_config dict (when config_or_params is params dict).

        skip_gp: bool flag to skip GP evaluation.

        skip_rm: bool flag to skip RM evaluation.

        

    Returns:

        dict with keys: 'flat', 'norm', 'planet', 'poly', 'fourier', 'trend',
@@ -58,6 +74,25 @@ def compute_model_components(params, plData, plModel, param_config, skip_gp=Fals 
        this is the flux baseline (near 1.0 for multiplicative mode, near 0.0 for

        additive mode). For RV instruments, this is the systemic velocity (gamma) [m/s].

    """

    if isinstance(config_or_params, ModelComponentsConfig):

        config = config_or_params

    else:

        config = ModelComponentsConfig(

            params=config_or_params,

            plData=plData,

            plModel=plModel,

            param_config=param_config,

            skip_gp=skip_gp,

            skip_rm=skip_rm,

        )



    params = config.params

    plData = config.plData

    plModel = config.plModel

    param_config = config.param_config

    skip_gp = config.skip_gp

    skip_rm = config.skip_rm



    time = np.array(plData.time)

    y = np.array(plData.y)

    dy = np.array(plData.dy)

PyAETNA/reporting/noise_plots.py

+9 −2
Changes for PyAETNA/reporting/noise_plots.py: 9 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -16,7 +16,7 @@ def _log(msg, level='INFO'): 
    logger._log(msg, level=level)



from .model_math import _apply_scoped_values

from .model_components import compute_model_components

from .model_components import compute_model_components, ModelComponentsConfig





def plot_noise_budget(plData, plModel, best_params, param_config,
@@ -88,7 +88,14 @@ def plot_noise_budget(plData, plModel, best_params, param_config, 
    if residuals_arr is not None:

        residuals = np.asarray(residuals_arr)

    else:

        model_comp = compute_model_components(best_params, plData, plModel, param_config)

        model_comp = compute_model_components(

            ModelComponentsConfig(

                params=best_params,

                plData=plData,

                plModel=plModel,

                param_config=param_config

            )

        )

        residuals = np.asarray(plData.y) - np.asarray(model_comp['total'])



    # -------------------------------------------------------- configuration

PyAETNA/reporting/pdf_report.py

+9 −2
Changes for PyAETNA/reporting/pdf_report.py: 9 added lines, 2 removed lines.
Original line number Diff line number Diff line
@@ -17,7 +17,7 @@ from .corner_plot import plot_corner, plot_kde_trace, plot_chain_comparison, plo 
from .plot_models import plot_models

from .plot_data_prep import prepare_plot_data

from .panel_nontime_residuals import plot_nontime_residuals

from .model_components import compute_model_components, save_model_components_binary

from .model_components import compute_model_components, save_model_components_binary, ModelComponentsConfig

from .confidence_bands import compute_data_bands

from .model_math import _precompute_sample_data

from .noise_plots import plot_noise_budget
@@ -218,7 +218,14 @@ def generate_report(pickle_path, output_pdf, log_file=None, config_file=None): 
        # Using the total model for residuals is critical: passing y - det_model

        # when a GP is active would leave the GP-absorbed signal in the residuals,

        # causing chi2_red >> 1 even for sigma_tot (which is expected to be ~1).

        model_comp = compute_model_components(map_params_full, plData, plModel, param_config)

        model_comp = compute_model_components(

            ModelComponentsConfig(

                params=map_params_full,

                plData=plData,

                plModel=plModel,

                param_config=param_config

            )

        )

        residuals_all = np.asarray(plData.y) - np.asarray(model_comp['total'])

        dy_tot_all    = np.asarray(model_comp['dy_total'])   # sigma_tot = sqrt(sigma_orig² + jitter²)

        dy_raw_all    = np.asarray(plData.dy)                # sigma_orig from the SIGMA column

PyAETNA/reporting/phase_plots.py

+35 −7
Changes for PyAETNA/reporting/phase_plots.py: 35 added lines, 7 removed lines.
Original line number Diff line number Diff line
@@ -49,7 +49,7 @@ def _log(msg, level='INFO'): 
from .gp_plotting import _get_gp_correlation_length, _precompute_gp_grid_context, _gp_predict_on_grid, _build_visit_time_grid, _build_lc_train_grid

from .corner_plot import _get_distinct_colors

from .model_math import _compute_planet_model_for_plotting, _apply_scoped_values

from .model_components import compute_model_components

from .model_components import compute_model_components, ModelComponentsConfig





def _time_to_phase(time, t0, period):
@@ -125,7 +125,14 @@ def plot_phase_folded(plData, plModel, best_params, param_config, planet_idx, 
        rp_rs = best_params[f'rp_by_rs_{planet_suffix}']



    # --- Compute model components once for the best-fit --------------------

    model_comp = compute_model_components(best_params, plData, plModel, param_config)

    model_comp = compute_model_components(

        ModelComponentsConfig(

            params=best_params,

            plData=plData,

            plModel=plModel,

            param_config=param_config

        )

    )

    time_full = np.array(plData.time)

    flux_full = np.array(plData.y)

    dy_full = np.array(plData.dy)
@@ -239,8 +246,15 @@ def plot_phase_folded(plData, plModel, best_params, param_config, planet_idx, 
        for si in idx_sel:

            p_i = {name: samples_dict[name][si] for name in samples_dict}

            p_i.update(param_config['fixed_params'])

            comp_i = compute_model_components(p_i, plData, plModel, param_config,

                                              skip_gp=True)  # only planet_terms needed

            comp_i = compute_model_components(

                ModelComponentsConfig(

                    params=p_i,

                    plData=plData,

                    plModel=plModel,

                    param_config=param_config,

                    skip_gp=True

                )

            )  # only planet_terms needed

            pt_i = comp_i.get('planet_terms', {})

            curve_i = pt_i.get(this_planet_key, np.zeros(len(time_full)))

            model_phase_arr.append(1.0 + curve_i[lc_mask])
@@ -279,8 +293,15 @@ def plot_phase_folded(plData, plModel, best_params, param_config, planet_idx, 
    )

    

    # Compute model on regular grid for smooth plotting

    comp_grid = compute_model_components(best_params, plData_grid, plModel, param_config,

                                         skip_gp=True)  # only planet_terms needed

    comp_grid = compute_model_components(

        ModelComponentsConfig(

            params=best_params,

            plData=plData_grid,

            plModel=plModel,

            param_config=param_config,

            skip_gp=True

        )

    )  # only planet_terms needed

    pt_grid = comp_grid.get('planet_terms', {})

    this_planet_grid = pt_grid.get(this_planet_key, np.zeros(n_grid))

    model_curve_grid = 1.0 + this_planet_grid  # smooth model on regular grid
@@ -659,7 +680,14 @@ def plot_phase_folded_rv(plData, plModel, best_params, param_config, planet_idx, 
    t0_rel = t0_abs - getattr(plData, 'timeCenter', 0.0)



    # Compute all model components

    model_comp = compute_model_components(best_params, plData, plModel, param_config)

    model_comp = compute_model_components(

        ModelComponentsConfig(

            params=best_params,

            plData=plData,

            plModel=plModel,

            param_config=param_config

        )

    )

    time = np.array(plData.time)

    flux = np.array(plData.y)  # This is RV in m/s for RV instruments

    dy = np.array(plData.dy)