feat: implement structured binary export for model components

# Jupyter

.ipynb_checkpoints/

*.ipynb

# Temporary files

*.tmp

import numpy as np

import jax

import jax.numpy as jnp

import dill 

from PyAETNA.gp import (

    gp_predict,

        _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.

    Args:

        outfile: Path to output .pckl file.

        map_data: Dict of components for MAP parameters.

        median_data: Dict of components for Median parameters.

        quantile_data: Dict of {name: array(n_pts, 2)} for 16th and 84th percentiles.

        plData: PlData object.

        plModel: PlModel object.

        param_config: Parameter configuration.

    """

    # 1. Basic Maps

    instrument_map = np.array(plData.instrument)

    visit_map = np.array(plData.visit)

    # 2. Instrument Metadata & Formulas

    instruments = {}

    data_types = param_config.get('data_types', getattr(plModel, 'data_types', None) or {})

    norm_mode = getattr(plModel, 'normalization_mode', 'multiplicative')

    inst_names_list = plModel.instrument_names if hasattr(plModel, 'instrument_names') and plModel.instrument_names is not None else None

    for i, inst_id in enumerate(plModel.uniqueInstruments):

        name = inst_names_list[i] if (inst_names_list is not None and i < len(inst_names_list)) else str(inst_id)

        dtype = data_types.get(inst_id, 'LC')

        # Construct a descriptive formula

        if dtype == 'LC':

            if norm_mode == 'multiplicative':

                formula = "Total = Norm * (1 + Planet + Poly + Fourier + EIV) + GP"

            else:

                formula = "Total = Norm + Planet + Poly + Fourier + EIV + GP"

        else: # RV

            formula = "Total = Norm + Planet + Poly + Fourier + EIV + GP"

        instruments[inst_id] = {

            'name': name,

            'formula': formula,

            'type': dtype

        }

    # 3. Visit Metadata

    visits = {}

    for uv in plModel.uniqueVisits:

        # We don't explicitly list every single term here to avoid bloat, 

        # but we can note the visit ID.

        visits[uv] = {'visit_id': uv}

    # 4. Component Names

    # Use keys from map_data as the master list

    all_component_keys = list(map_data.keys())

    # Assemble final structure

    export_dict = {

        'metadata': {

            'instrument_map': instrument_map,

            'visit_map': visit_map,

            'instruments': instruments,

            'visits': visits,

            'component_names': all_component_keys,

        },

        'coords': {

            'bjd': map_data.get('bjd', np.array(plData.time) + getattr(plData, 'timeCenter', 0.0)),

        },

        'map_values': map_data,

        'median_values': median_data,

        'quantiles': quantile_data,

    }

    try:

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

            dill.dump(export_dict, f)

        _log(f"Model components binary export saved to: {outfile}")

    except Exception as e:

        _log(f"ERROR: failed to save binary components: {str(e)}", level='ERROR')

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, write_model_components

from .model_components import compute_model_components, write_model_components, save_model_components_binary

from .confidence_bands import compute_data_bands

from .model_math import _precompute_sample_data

from .noise_plots import plot_noise_budget

from .periodogram_plots import (

    _log(f"Report saved: {output_pdf}")

    # Write overall model components file (full dataset, not per-visit or per-instrument subsets)

    # Write overall model components binary file (MAP, Median, Quantiles)

    # Data are NOT sorted by time to maintain 1-to-1 correspondence with the original dataset

    try:

        components_file = output_pdf.replace('.pdf', '_components.txt')

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

        write_model_components(components_file, model_comp, plData=plData, mask=None, sort_by_time=False)

        _log(f"  - Overall model components saved: {components_file}")

        components_bin = output_pdf.replace('.pdf', '_components.pckl')

        # 1. MAP values

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

        # 2. Median and Quantiles from confidence_bands

        # We need to prepare the arguments for compute_data_bands

        # Note: compute_data_bands expects samples_dict, sample_cache, etc.

        # We use the samples from the loaded pickle.

        # Prepare basic config for bands

        has_planets = len(param_config.get('planet_names', ())) > 0

        has_non_time_trends = len(param_config.get('polynomial_configs', [])) > 0 or \

                              len(param_config.get('fourier_configs', [])) > 0 or \

                              len(param_config.get('eiv_configs', [])) > 0

        # Extract necessary configs

        poly_cfgs = param_config.get('polynomial_configs', [])

        four_cfgs = param_config.get('fourier_configs', [])

        # Call compute_data_bands to get posterior statistics

        bands = compute_data_bands(

            samples_dict=samples,

            sample_cache=None,

            inst=np.array(plData.instrument),

            data_types_config=param_config.get('data_types', getattr(plModel, 'data_types', None) or {}),

            plData=plData,

            plModel=plModel,

            param_config=param_config,

            normalization_mode=getattr(plModel, 'normalization_mode', 'multiplicative'),

            has_planets=has_planets,

            has_non_time_trends=has_non_time_trends,

            poly_cfgs=poly_cfgs,

            four_cfgs=four_cfgs,

        )

        # Map bands output to our binary structure

        # compute_data_bands returns keys like 'total_median', 'total_q16', etc.

        median_values = {}

        quantile_values = {}

        # We extract the components that are actually computed.

        # For the total model:

        median_values['total'] = bands['total_median']

        quantile_values['total'] = np.stack([bands['total_q16'], bands['total_q84']], axis=-1)

        # For the mean trend:

        median_values['mean_function'] = bands['mean_trend_median']

        quantile_values['mean_function'] = np.stack([bands['mean_trend_q16'], bands['mean_trend_q84']], axis=-1)

        # For planets (if any):

        if has_planets:

            median_values['planet'] = bands['planet_median']

            quantile_values['planet'] = np.stack([bands['planet_q16'], bands['planet_q84']], axis=-1)

            median_values['norm'] = bands['planet_norm_median']

            quantile_values['norm'] = np.stack([bands['planet_norm_q16'], bands['planet_norm_q84']], axis=-1)

            median_values['trend'] = bands['planet_trend_median']

            quantile_values['trend'] = np.stack([bands['planet_trend_q16'], bands['planet_trend_q84']], axis=-1)

        # For non-time trends (EIV, etc.)

        if has_non_time_trends and bands.get('nt_data_q16') is not None:

            nt_q16 = bands['nt_data_q16']

            nt_q50 = bands['nt_data_q50']

            nt_q84 = bands['nt_data_q84']

            for name in nt_q16.keys():

                median_values[name] = nt_q50[name]

                quantile_values[name] = np.stack([nt_q16[name], nt_q84[name]], axis=-1)

        # Save to binary

        save_model_components_binary(

            components_bin, map_comp, median_values, quantile_values, 

            plData, plModel, param_config

        )

    except Exception as e:

        _log(f"WARNING: failed to write overall model components: {str(e)}")

        _log(f"WARNING: failed to write overall model components binary: {str(e)}")

        import traceback

        _log(traceback.format_exc())

    # Save MAP result back to pickle so planet_phys.py can reuse it

    try:

        data['map_result'] = map_result

        with open(pickle_path, 'wb') as f:

            dill.dump(data, f)

        _log(f"MAP result cached in pickle: {pickle_path}")

        _log(f"Full session state (including MAP result) cached in pickle: {pickle_path}")

    except Exception as e:

        _log(f"WARNING: failed to save MAP result to pickle: {e}")

  - Version compatibility and troubleshooting

  - Conda and Docker alternatives

- **Structured Model Exports.** In addition to visual reports saved as a multi-page pdf file, the pipeline exports the full model (`*.pckl`) state and a dedicated binary components file (`*_components.pckl`). This allows users to programmatically access the best-fit values and posterior uncertainties (16-84% quantiles) for every individual model component at every data point.

Additional resources:

- Test files in `tests/`

- Inline documentation in source code

## Pipeline Outputs

After a successful run, PyAETNA generates several output files in the output directory:

- **Logfile (`.log`)**: A detailed textual record of the entire process, including configuration, sampling progress, convergence diagnostics, and a tailored list of citations.

- **PDF Report (`.pdf`)**: A comprehensive visual summary of the fit, containing the plots described in the "Features" section.

- **Session Cache (`.pckl`)**: A binary snapshot of the full MCMC state. This file is used for checkpointing and allows you to resume a run from where it left off.

- **Model Components Export (`_components.pckl`)**: A structured binary file containing:

    - `coords`: Time coordinates (BJD).

    - `map_values`: Best-fit values for each model component.

    - `median_values`: Posterior median values.

    - `quantiles`: 16% and 84% posterior quantiles.

    - `metadata`: Instrument maps and the formulas used to reconstruct the total model.

For a practical example of how to load and use this binary export to reconstruct your model, see the [Load Binary Model Example](docs/load_binary_model.ipynb).

## Repository Structure

```