refactor of group 3 (sampling): complete Group 3 code review for data loading,...

"""Observational data containers and column-validation utilities for PyAETNA.

The :class:`PlData` class is the central data container used throughout the

sampling engine.  Raw ASCII tables are first validated and normalised by

:func:`_validate_and_prepare_columns` before being handed to ``PlData``.

"""

from __future__ import annotations

from typing import Callable, Optional

import numpy as np

import jax.numpy as jnp

import numpy.lib.recfunctions as rfn

_SYSTEM_COLUMNS = {'BJD', 'Y', 'SIGMA', 'EXPTIME', 'VISIT', 'INSTRUMENT'}

def _validate_and_prepare_columns(data, log_fn=None):

def _validate_and_prepare_columns(

    data: np.ndarray,

    log_fn: Optional[Callable[[str], None]] = None,

) -> np.ndarray:

    """Validate mandatory data columns and fill in defaults for optional ones.

    Mandatory (error and exit if absent):

    columns before constructing a ``PlData`` object.

    """

    def __init__(self, data):

    def __init__(self, data: np.ndarray) -> None:

        self.time = data['BJD']

        self.timeCenter = (self.time.max() + self.time.min()) / 2  # Store for tau adjustment

        # window to account for integration-time smearing (Kipping 2010).

        # Typically constant within each visit; varying within a visit

        # triggers a warning in _build_numpyro_model.

        if 'EXPTIME' in data.dtype.names:

        # NOTE: _validate_and_prepare_columns guarantees EXPTIME is always

        # present in *data* before PlData is constructed, so no else-branch

        # is needed here.

        self.exptime = jnp.array(data['EXPTIME'])

        else:

            self.exptime = jnp.zeros(len(self.y))

        self.visit = jnp.array(data['VISIT'].astype('int'))

        self.uniqueVisits = np.unique(data['VISIT'].astype('int'))

            self.dBasis = {name: jnp.array(basis_data[name]) for name in basis_names}

        else:

            self.dBasis = {}

    def __repr__(self) -> str:

        n_pts = int(len(self.time))

        n_visits = int(len(self.uniqueVisits))

        n_instr = int(len(jnp.unique(self.instrument)))

        basis_cols = list(self.dBasis.keys())

        basis_str = f", dBasis={basis_cols}" if basis_cols else ""

        return (

            f"PlData(n={n_pts}, visits={n_visits}, instruments={n_instr}"

            f"{basis_str})"

        )

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

"""Checkpoint loading, saving, and precision upcasting/downcasting utilities."""

"""Checkpoint loading, saving, and precision upcasting/downcasting utilities.

Checkpoints are serialised with :mod:`dill` so that arbitrary Python objects

(including NumPyro/JAX MCMC states) can be persisted across runs.  On reload

the stored precision (float32/float64) is reconciled with the current

configuration via :func:`_configure_jax_precision`.

"""

from __future__ import annotations

from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import os

import dill

import numpy as np

import jax.numpy as jnp

from ..precision import _configure_jax_precision, _cast_warmup_state_dtype

def upcast_value_to_float64(val):

    """Recursively upcast JAX arrays and nested structures to float64/int64 (x64 mode)"""

def upcast_value_to_float64(val: Any) -> Any:

    """Recursively upcast JAX arrays and nested structures to float64/int64 (x64 mode).

    Supported container types: ``dict``, ``list``, named-tuples (``_replace``

    protocol), and any object with a ``dtype`` + ``shape`` attribute (JAX/NumPy

    arrays).  Unrecognised types are returned unchanged.

    .. note::

        This function is a no-op (returns *val* unchanged) when JAX x64 mode is

        disabled (i.e. ``jax.config.jax_enable_x64`` is ``False``), because JAX

        would silently truncate the requested dtype back to float32/int32 anyway.

        Enable x64 mode via ``jax.config.update('jax_enable_x64', True)`` *before*

        calling this function.

    """

    import jax

    if not jax.config.jax_enable_x64:

        # x64 unavailable – casting would silently truncate; return as-is.

        return val

    if isinstance(val, dict):

        return {k: upcast_value_to_float64(v) for k, v in val.items()}

    elif isinstance(val, list):

        # Recurse into lists so that lists of JAX arrays are upcast correctly.

        return [upcast_value_to_float64(v) for v in val]

    elif hasattr(val, 'dtype') and hasattr(val, 'shape'):

        if val.dtype == jnp.float32:

            return jnp.array(val, dtype=jnp.float64)

    else:

        return val

def save_pickle(save_path, samples_dict, plModel, plData, config, param_config,

                num_chains, requested_num_chains, num_warmup, device_type,

                sampler_type, mcmc_last_state=None, extra=None, verbose=False, log_func=print):

    """Save samples + state to disk."""

def save_pickle(

    save_path: str,

    samples_dict: Dict[str, Any],

    plModel: Any,

    plData: Any,

    config: Dict[str, Any],

    param_config: Dict[str, Any],

    num_chains: int,

    requested_num_chains: int,

    num_warmup: int,

    device_type: str,

    sampler_type: str,

    mcmc_last_state: Optional[Any] = None,

    extra: Optional[Dict[str, Any]] = None,

    verbose: bool = False,

    log_func: Callable[..., None] = print,

) -> None:

    """Serialise samples and sampler state to *save_path* with :mod:`dill`.

    Parameters

    ----------

    verbose:

        When *True*, emit a log line before writing.  Other progress messages

        (e.g. from the calling context) are not suppressed by this flag.

    extra:

        Optional extra keys merged into the checkpoint dict (e.g.

        ``warmup_mcmc_state``, ``warmup_done``).

    """

    if verbose:

        log_func(plModel, f"Saving to {save_path}...")

    min_flat = min((len(v) for v in samples_dict.values()), default=0)

    num_samples_stored = min_flat // num_chains if num_chains else 0

    # Store per-parameter leading dimension (chain count) so that

    # load_checkpoint can recover num_chains reliably without heuristics.

    samples_shapes = {

        k: list(v.shape) for k, v in samples_dict.items() if hasattr(v, 'shape')

    }

    out = {

        'plModel': plModel,

        'plData': plData,

        'config': config,

        'param_config': param_config,

        'samples': samples_dict,

        'samples_shapes': samples_shapes,

        'num_chains': num_chains,

        'requested_num_chains': requested_num_chains,

        'num_samples': num_samples_stored,

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

        dill.dump(out, f, protocol=4)

def load_checkpoint(save_path, reset, plModel, config, mcmc_config_peek, log_func, format_time, logFile):

    """Load existing checkpoint and restore sampler states, adapting precision."""

def load_checkpoint(

    save_path: str,

    reset: bool,

    plModel: Any,

    config: Dict[str, Any],

    mcmc_config_peek: Dict[str, Any],

    log_func: Callable[..., None],

    format_time: Callable[[float], str],

    logFile: Any,

) -> Tuple[

    Optional[Dict[str, Any]],  # existing_samples

    int,                        # existing_num_samples

    bool,                       # loaded_from_pickle

    str,                        # saved_sampler_type

    Optional[Any],              # existing_mcmc_state

    Optional[Any],              # existing_warmup_state

    int,                        # existing_warmup_done

    float,                      # existing_warmup_elapsed

    Optional[int],              # num_chains

    Any,                        # plModel

    Optional[Any],              # plData

    Dict[str, Any],             # config

    Optional[Dict[str, Any]],   # param_config

    bool,                       # use_float64

]:

    """Load an existing checkpoint and restore sampler states, adapting precision."""

    plData = None

    existing_samples = None

    existing_num_samples = 0

                        }

            existing_samples = saved_data.get('samples', None)

            if not existing_samples:

            # Explicit None/empty check: using `not existing_samples` would

            # raise ValueError if existing_samples were a numpy array.

            if existing_samples is not None and len(existing_samples) == 0:

                existing_samples = None

            existing_mcmc_state = saved_data.get('mcmc_last_state', None)

            if existing_samples is not None:

                scalar_flat = min(len(v) for v in existing_samples.values())

                actual_chain_count = None

                # Priority 1: use per-parameter shape stored by save_pickle

                # (available in checkpoints written after this fix).

                samples_shapes = saved_data.get('samples_shapes', {})

                if samples_shapes:

                    first_shape = next(iter(samples_shapes.values()), None)

                    if first_shape and len(first_shape) >= 2:

                        # Shape is (chains * samples,) flat or (chains, samples, ...)

                        # save_pickle stores the flat shape, so leading dim = chains

                        # We stored shapes BEFORE flattening: skip if already flat

                        actual_chain_count = first_shape[0]

                        log_func(logFile, f"Inferred {actual_chain_count} chains from stored samples_shapes")

                # Priority 2: inspect MCMC state z-array leading dimension

                if actual_chain_count is None:

                    if existing_mcmc_state is not None and hasattr(existing_mcmc_state, 'z'):

                        first_walker_param = next(iter(existing_mcmc_state.z.values()), None)

                        if first_walker_param is not None:

                            actual_chain_count = len(first_walker_param)

                            log_func(logFile, f"Inferred {actual_chain_count} chains from MCMC state z shape")

                # Priority 3: heuristic divisibility fallback (legacy checkpoints)

                if actual_chain_count is None:

                    # Heuristic fallback: try common chain counts used in

                    # PyAETNA runs (4 for quick tests, 16/32 for production,

                    # 34 for legacy runs) and pick the first that divides

                    # the flat sample count evenly.

                    candidate_counts = [4, 16, 32, 34]

                    for candidate in candidate_counts:

                        if scalar_flat % candidate == 0:

                            actual_chain_count = candidate

                            log_func(logFile, f"Inferred {actual_chain_count} chains from scalar param divisibility")

                            log_func(logFile, f"Inferred {actual_chain_count} chains from scalar param divisibility (heuristic)")

                            break

                if actual_chain_count is None:

                    log_func(logFile, f"Using saved num_chains={actual_chain_count} (may be incorrect if production used different sampler)")

                num_chains = actual_chain_count

                existing_num_samples = scalar_flat // num_chains

                existing_num_samples = scalar_flat // num_chains if num_chains else 0

                log_func(logFile, f"Loaded checkpoint: {len(plData.time)} data points, "

                         f"{existing_num_samples} samples/chain × {num_chains} chains "

            except Exception:

                plData.timeZeroPoint = getattr(plData, 'timeZeroPoint', 0.0)

            mcmc_config_restored = config.get('mcmc', {})

            mcmc_config_restored.update({

                k: v for k, v in mcmc_config_peek.items()

                if k == 'use_float64'

            })

            use_float64 = _configure_jax_precision(param_config, mcmc_config_restored, log_func=lambda *args: log_func(logFile, args[-1]))            

            # Build a *new* dict (do not mutate the loaded config) so that

            # the on-disk config remains unchanged if the caller inspects it.

            mcmc_config_restored = {**config.get('mcmc', {}),

                                    **{k: v for k, v in mcmc_config_peek.items()

                                       if k == 'use_float64'}}

            use_float64 = _configure_jax_precision(

                param_config,

                mcmc_config_restored,

                log_func=lambda *args: log_func(logFile, args[-1]),

            )            

            if use_float64:

                for attr in ('time', 'y', 'dy', 'exptime'):

                    arr = getattr(plData, attr, None)

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

"""Data loading, validation, and preprocessing utilities."""

"""Data loading, validation, and preprocessing utilities.

This module orchestrates the full data + configuration pipeline:

file I/O → column validation → :class:`~PyAETNA.sampling.data.PlData`

construction → YAML config parsing → precision setup → GP kernel validation.

Configuration errors are raised as :exc:`~PyAETNA.sampling.exceptions.PyAETNAConfigError`

so that the CLI entry-point can catch them, log a user-friendly message, and

exit cleanly.  Direct ``sys.exit`` calls have been removed from library code.

"""

from __future__ import annotations

from typing import Any, Callable, Dict, List, Tuple

import os

import sys

import numpy as np

from ..config.parser import parse_config

from ..precision import _configure_jax_precision

from .logging import _log_config_summary

def load_data_and_config(dataFile, cfg_path, config, plModel, log_func):

from .exceptions import PyAETNAConfigError

from PyAETNA.kernels.factory import SUPPORTED_KERNELS, _count_spectral_mixture_components

def load_data_and_config(

    dataFile: str,

    cfg_path: str,

    config: Dict[str, Any],

    plModel: Any,

    log_func: Callable[..., None],

) -> Tuple[PlData, Any, Dict[str, Any], bool]:

    """Load, validate, and preprocess the dataset and configuration."""

    log_func(plModel, f'Loading data from {dataFile}...')

    plModel.inst_masks = plModel.instrument_masks

    visit_to_inst_id = []

    for v_mask in plModel.visit_masks:

    for v_idx, v_mask in enumerate(plModel.visit_masks):

        v_inst_id = None

        for j, i_mask in enumerate(plModel.instrument_masks):

            if np.any(np.logical_and(np.array(v_mask), np.array(i_mask))):

                v_inst_id = plModel.uniqueInstruments[j]

                break

        if v_inst_id is None:

            log_func(

                plModel,

                f"Warning: visit index {v_idx} (visit mask #{v_idx}) does not "

                "overlap with any instrument mask — visit_to_inst_id entry will "

                "be None, which may cause downstream errors.",

            )

        visit_to_inst_id.append(v_inst_id)

    plModel.visit_to_inst_id = visit_to_inst_id

    if 'instruments' in config and len(plModel.uniqueInstruments) != len(instrument_names):

        error_msg = f"Mismatch: data has {len(plModel.uniqueInstruments)} unique instrument(s) but YAML 'instruments' contains {len(instrument_names)} entries. Check INSTRUMENT column and YAML 'instruments' list."

        log_func(plModel, f"ERROR: {error_msg}")

        log_func(plModel, "Aborting run.")

        sys.exit(1)

        raise PyAETNAConfigError(error_msg)

    data_type_list = config.get('data_type', None)

    if data_type_list is None:

            f"Valid values: 'LC' (light curve) or 'RV' (radial velocity)."

        )

        log_func(plModel, f"ERROR: {error_msg}")

        log_func(plModel, "Aborting run.")

        sys.exit(1)

        raise PyAETNAConfigError(error_msg)

    if len(data_type_list) != len(instrument_names):

        error_msg = (

            f"'data_type' list length ({len(data_type_list)}) must match "

            f"Provide one entry ('LC' or 'RV') per instrument."

        )

        log_func(plModel, f"ERROR: {error_msg}")

        log_func(plModel, "Aborting run.")

        sys.exit(1)

        raise PyAETNAConfigError(error_msg)

    for dt in data_type_list:

        if dt.upper() not in ('LC', 'RV'):

            error_msg = f"Invalid data_type '{dt}'. Must be 'LC' (light curve) or 'RV' (radial velocity)."

            log_func(plModel, f"ERROR: {error_msg}")

            log_func(plModel, "Aborting run.")

            sys.exit(1)

            raise PyAETNAConfigError(error_msg)

    plModel.data_types = {

        inst_id: dt.upper()

        for inst_id, dt in zip(plModel.uniqueInstruments, data_type_list)

        plData.dBasis['ROLL_ANGLE'] = jnp.array(temp)

        if not roll_unfold:

            dRA = np.diff(np.array(plData.dBasis['ROLL_ANGLE']))

            roll_np = np.array(plData.dBasis['ROLL_ANGLE'])  # explicit conversion for np.diff

            dRA = np.diff(roll_np)

            maxD = np.argmax(dRA)

            zeroPoint = 0.5 * (plData.dBasis['ROLL_ANGLE'][maxD] + plData.dBasis['ROLL_ANGLE'][maxD+1])

            zeroPoint = 0.5 * (roll_np[maxD] + roll_np[maxD + 1])

            plData.dBasis['ROLL_ANGLE'] = plData.dBasis['ROLL_ANGLE'] - zeroPoint

            negRA = jnp.where(plData.dBasis['ROLL_ANGLE'] < 0)[0]

            plData.dBasis['ROLL_ANGLE'] = plData.dBasis['ROLL_ANGLE'].at[negRA].set(

                f"but column '{_col}' was not found in data file. "

                f"Available extra columns: {list(plData.dBasis.keys())}. Aborting."

            ))

            sys.exit(1)

            raise PyAETNAConfigError(

                f"EIV block '{_eiv_name}': x_variable column '{_col}' not found in data file."

            )

        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)

            raise PyAETNAConfigError(

                f"EIV block '{_eiv_name}': uncertainty column '{_err_col}' not found in data file."

            )

        log_func(plModel, f"EIV block '{_eiv_name}': validated columns '{_col}' and '{_err_col}'. OK.")

    try:

    except ValueError as e:

        error_msg = f"Configuration error: {str(e)}"

        log_func(plModel, f"ERROR: {error_msg}")

        log_func(plModel, "Aborting run.")

        sys.exit(1)

        raise PyAETNAConfigError(error_msg) from e

    mcmc_config_early = config.get('mcmc', {})

    for inst_id, gp_cfg in gp_configs.items():

        for spec in gp_cfg.get('kernel_specs', []):

            kernel_type = spec.get('type', 'Matern32')

            from PyAETNA.kernels.factory import SUPPORTED_KERNELS

            # SUPPORTED_KERNELS and _count_spectral_mixture_components are

            # imported at module level to avoid repeated import resolution.

            if kernel_type.lower() not in [k.lower() for k in SUPPORTED_KERNELS]:

                msg = f"Unknown GP kernel type '{kernel_type}'. Supported: {', '.join(SUPPORTED_KERNELS)}"

                raise ValueError(msg)

            if kernel_type.lower() == 'spectralmixture':

                try:

                    from PyAETNA.kernels.factory import _count_spectral_mixture_components

                    param_names = spec.get('params', [])

                    fixed_params = spec.get('fixed_params', {})

                    per_param_scopes = spec.get('param_scopes', {})

                    default_scope = spec.get('scope', 'per_instrument')

                        inst_suffix=None,

                        section_name=inst_id,

                        per_param_scopes=per_param_scopes,

                        default_scope=default_scope

                        default_scope=default_scope,

                    )

                    spec['cached_component_count'] = component_count

                except Exception as e:

                    raise ValueError(f"Failed to compute component count for SpectralMixture kernel in section '{inst_id}': {e}")

                    raise ValueError(

                        f"Failed to compute component count for SpectralMixture "

                        f"kernel in section '{inst_id}': {e}"

                    )

    _log_config_summary(param_config, plModel, plData, instrument_names, log_func)

    #raise RuntimeError("Priors parsed and printed. Stopping execution as requested.")

    plModel.param_config = param_config

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

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

"""Custom exception hierarchy for the PyAETNA sampling engine.

Raising typed exceptions (rather than calling ``sys.exit``) from library code

allows callers \u2014 including tests and higher-level orchestrators \u2014 to catch and

handle errors programmatically.  The CLI entry-point in ``run.py`` catches these

exceptions and exits with an informative message.

"""

from __future__ import annotations

class PyAETNAError(Exception):

    """Base class for all PyAETNA-specific exceptions."""

class PyAETNAConfigError(PyAETNAError):

    """Raised when the YAML configuration is invalid or inconsistent with the data.

    Examples include missing mandatory keys, mismatched list lengths, and

    unknown ``data_type`` values.

    Parameters

    ----------

    message:

        Human-readable description of the configuration problem.

    """

    def __init__(self, message: str) -> None:

        super().__init__(message)

        self.message = message

    def __str__(self) -> str:  # pragma: no cover

        return f"Configuration error: {self.message}"

                    # Use pre-calculated mapping to avoid TracerBoolConversionError.

                    if visit_to_inst_id is not None:

                        inst_id = visit_to_inst_id[i]

                        if inst_id is None:

                            raise ValueError(

                                f"Polynomial '{poly_name}': visit index {i} (visit {uv}) "

                                f"has no associated instrument in visit_to_inst_id. "

                                f"Check that every visit overlaps with exactly one instrument mask."

                            )

                    else:

                        # Fallback (may fail in JIT context)

                        inst_id = None