refactor of group 7 (logging and error handling) started

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

from .logger import set_log_file, _get_logger, _log, _LOG_FILE

# NOTE: `_log`/`_LOG_FILE` are intentionally NOT re-exported here. Both are

# rebound at runtime by `set_log_file()`, so importing them by value (e.g.

# `from PyAETNA.reporting import _log`) would silently capture a stale

# reference. Use `from PyAETNA.reporting import logger` and call

# `logger._log(...)`, or call `_get_logger(...)` for a fresh instance.

from .logger import set_log_file, _get_logger

from .map_search import (

    find_map_parameters,

    _evaluate_chain_endpoints,

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

"""

Module-level logging utilities shared by the PyAETNA reporting package.

Provides a small file-or-stdout logger (`_get_logger`) plus a module-wide

default instance (`_log`) that can be redirected at runtime via

`set_log_file`.

"""

import time as ttt

import sys

_LOG_FILE = None

_log = None

    def _log(msg, level='INFO'):

        timestamp = ttt.strftime("%Y-%m-%d %H:%M:%S", ttt.gmtime())

        formatted = f"{timestamp} {level}: {msg}"

        try:

            if log_file:

                with open(log_file, 'a') as f:

                    f.write(formatted + '\n')

            else:

        if not log_file:

            print(formatted)

        except Exception:

            # Best-effort: fallback to module-level log file if set, otherwise stdout

            return

        try:

                # Use module-level log file if available

                global _LOG_FILE

                if _LOG_FILE:

                    with open(_LOG_FILE, 'a') as f:

            with open(log_file, 'a') as f:

                f.write(formatted + '\n')

                else:

                    print(formatted)

            except Exception:

                try:

                    print(formatted)

                except Exception:

                    pass

        except OSError as e:

            # Fall back to stdout, but surface the failure instead of

            # swallowing it silently so a broken log path is noticeable.

            print(f"{formatted} (warning: failed to write to log file {log_file!r}: {e})")

    return _log

from ..utils import linearizeRollAngle, relabel_array

from ..config.parser import parse_config

from ..precision import _configure_jax_precision

from .logging import _log_config_summary

from .sampling_logging import _log_config_summary

from .exceptions import PyAETNAConfigError

from PyAETNA.kernels.factory import SUPPORTED_KERNELS, _count_spectral_mixture_components

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

                    )

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

    _log_config_summary(param_config, plModel, instrument_names, log_func)

    plModel.param_config = param_config

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

from numpyro.infer import MCMC, NUTS, init_to_median

from .context import SamplingContext

from .logging import _format_rate, check_walker_initialization as _cwi

from .sampling_logging import _format_rate, check_walker_initialization as _cwi

from .phases import (

    run_warmup_phase,

    run_nuts_only_fallback,

            n_extra = n_total // n_scalar

            samples_grouped[name] = vals.reshape(num_chains, existing_num_samples, n_extra)

        else:

            _log(plModel, "WARNING: Checkpoint incompatible with current configuration.")

            _log(plModel, "Checkpoint incompatible with current configuration.", level='WARNING')

            _log(plModel, f"  Parameter '{name}': {n_total} values cannot be split into "

                 f"{num_chains} chains × {existing_num_samples} samples (or integer multiples).")

            _log(plModel, "  This usually means the dataset or free-parameter set has changed.")

        mcmc._warmup_state = existing_mcmc_state

        _log(plModel, "Restored saved MCMC state (adapted mass matrix & step size)")

    else:

        _log(plModel, "WARNING: No saved MCMC state. Starting from last sample positions (may diverge).")

        _log(plModel, "No saved MCMC state. Starting from last sample positions (may diverge).", level='WARNING')

        # Must run one batch with init_params, then continue from state

        mcmc.run(jax.random.PRNGKey(random_seed + existing_num_samples),

                 pl_data=plData, pl_model=plModel, init_params=init_params,

from numpyro.infer import MCMC, NUTS, init_to_median, init_to_value

from .context import SamplingContext

from .logging import _format_rate, _log_warmup_summary

from .sampling_logging import _format_rate, _log_warmup_summary

from .numpyro_model import _compute_log_joint

from .sampling_utils import (

    save_pickle,

                warmup_energy[chain_id, 0]  = -_compute_log_joint(model_fn, plData, plModel, p_first)

                warmup_energy[chain_id, -1] = -_compute_log_joint(model_fn, plData, plModel, p_last)

            except Exception as e:

                _log(plModel, f"WARNING: Failed to compute warmup energy for chain {chain_id}: {e}")

                _log(plModel, f"Failed to compute warmup energy for chain {chain_id}: {e}", level='WARNING')

                warmup_energy[chain_id, 0] = np.nan

                warmup_energy[chain_id, -1] = np.nan

        _log_warmup_summary(param_config, warmup_samples, warmup_energy, num_chains, plModel, _log)

    if has_nan:

        _log(plModel, "")

        _log(plModel, "WARNING: NaN detected in walker positions after ensemble warmup!")

        _log(plModel, "NaN detected in walker positions after ensemble warmup!", level='WARNING')

        _log(plModel, "This typically means the model produced NaN log-probabilities")

        _log(plModel, "during slice sampling (e.g., transit math with out-of-bounds params).")

        _log(plModel, "Walker positions (NaN count per param):")