added check_walker_initialization() to validate the initialization of walkers...

# Strategy: Eliminating Model Safeguards and Fixing Boundary Initialization

## Problem Description

PyAETNA currently employs several numerical safeguards (e.g., `jnp.clip`, `jnp.where`, and `jnp.maximum`) within its physical models (specifically in `body_models.py`) to prevent `NaN` values during the forward pass. While these prevent immediate crashes, they mask a systemic issue: **Boundary Initialization**.

As analyzed in `boundary_init_analysis.md`, when using the `auto` sampling pipeline with `init_to_uniform`, bounded priors (such as `TruncatedNormal`) are frequently initialized **exactly** at their lower or upper boundaries. 

### The Failure Mechanism:

1. **Unconstrained Space**: Samplers operate in unconstrained space $(-\infty, \infty)$. A value exactly at the boundary of a constrained prior maps to $\pm\infty$ in unconstrained space.

2. **Gradient Collapse**: At $\pm\infty$, the gradient of the transformation becomes zero or `NaN`.

3. **The "Invisible Wall"**: Model safeguards (like clipping) allow the forward pass to succeed, but the sampler remains stranded on a flat plateau with no gradient to guide it back into the high-probability region. This leads to numerical instability, extremely slow warmup, or total failure of the ensemble phase.

## Goal

Transition from a "safeguarded" model that hides errors to an "honest" model that prevents them by design. This involves moving safety from the **Model (Forward Pass)** to the **Prior/Initialization (Input Stage)**.

---

## Detailed Implementation Plan

### Step 1: Boundary Detection & Validation (The "Honesty" Check)

Before proceeding with the purge of safeguards, we must implement a strict validation step to detect boundary hits.

**Proposed Function: `validate_initial_walkers`**

This function should be called after the sampler has initialized the walkers but **before** the JIT-compiled sampling loop begins.

**Logic:**

1. Iterate through all initialized parameters.

2. For each parameter with a bounded prior $[low, high]$:

    - Check if any walker value $\theta_i$ satisfies $\theta_i = low$ or $\theta_i = high$ (within the numerical precision ised by the pipeline).

3. If boundary hits are detected:

    - Log the specific parameter name.

    - Log the index of the walker(s) hitting the boundary and the unconstrained `z`.

    - Log the exact value and the prior bounds.

    - **Raise a `RuntimeError` and exit.** The error must be logged to the same logfile used by the rest of the pipeline.

*This ensures that we cannot start a run that is mathematically doomed to fail due to $z = \pm\infty$.*

### Step 2: Hardening the Priors

To make the model safe without internal clipping, we shift the boundaries of the priors themselves.

1. **$\epsilon$-Shrinking**: Update the prior parsing logic to automatically shrink bounded priors by a small margin $\epsilon$ (e.g., $10^{-6}$).

   - $[low, high] \rightarrow [low + \epsilon, high - \epsilon]$.

2. **Result**: This ensures that even if a walker is at the very edge of the prior, it is still strictly within the valid domain of functions like `arccos` and `sqrt`.

### Step 3: Fixing the Initialization Root Cause

Address the bug in the interaction between `init_to_uniform` and bounded distributions.

1. **Custom Init Wrapper**: If `init_to_uniform` continues to produce boundary values, implement a custom initialization wrapper that draws samples from the interior of the support.

2. **Verification**: Use the `validate_initial_walkers` function from Step 1 to verify that the new initialization strategy produces zero boundary hits.

### Step 4: Purging the Safeguards

Once the input stage is guaranteed to be safe, remove the "crutches" from the physics engine.

1. **Remove `jnp.clip`**: Replace clipped calls (e.g., in `compute_inclination`) with direct calls.

2. **Remove `jnp.where` for safety**: Remove the `b_safe` logic in `compute_transit_duration` and `compute_eclipse_duration`.

3. **Remove `jnp.maximum(..., 0.0)`**: Remove the flooring of square root arguments.

## Expected Outcome

- **Numerical Honesty**: Any `NaN` encountered will be a legitimate signal of a model or prior failure, not a masked symptom.

- **Faster Convergence**: Walkers will start in the interior of the unconstrained space, providing healthy gradients for the sampler.

- **Cleaner Code**: The physics models in `body_models.py` will be stripped of boilerplate safety logic, making them easier to maintain and read.

"""

Logging utilities for the PyAETNA sampling engine.

"""

import jax

import jax.numpy as jnp

def _format_sampling_message(remaining, checkpoint_interval):

    """Return the user-friendly log string for a sampling batch.

    """

    return f"{rate:.3f}"

def check_walker_initialization(mcmc_obj, plModel, log_func, prob_fn=None, plData=None):

    """

    Log the initial walker positions in unconstrained space to verify initialization.

    Computes potential energy and gradients on-the-fly for each walker.

    Args:

        mcmc_obj: MCMC sampler object

        plModel: Model object (used for logging)

        log_func: Logging function

        prob_fn: Optional probability function to compute potential energy/gradients.

                If None, attempts to extract from mcmc_obj.kernel.

        plData: Data object required by prob_fn.

    """

    state = mcmc_obj.last_state

    if state is None:

        log_func(plModel, "  Warning: mcmc_obj.last_state is None. Cannot check initialization.")

        return

    z_values = getattr(state, 'z', getattr(state, 'position', None))

    log_func(plModel, "\n" + "="*150)

    log_func(plModel, "=== WALKER POSITIONS AFTER FIRST STEP ===")

    log_func(plModel, "="*150)

    if isinstance(z_values, dict):

        # Get first walker to see the structure

        try:

            first_key = next(iter(z_values.keys()))

            num_walkers = z_values[first_key].shape[0]

        except (StopIteration, AttributeError):

            log_func(plModel, "  z_values dict is empty or contains no arrays.")

            return

        # Get constrained samples for side-by-side comparison

        constrained_samples = None

        try:

            constrained_samples = mcmc_obj.get_samples(group_by_chain=True)

        except Exception:

            pass

        # Track statistics for warnings

        walker_stats = {i: {'has_nan': False, 'has_inf': False, 'max_z_bounded': None} for i in range(num_walkers)}

        # Identify which parameters are bounded (constrained)

        # In NumPyro, parameters with bounded priors are transformed to unconstrained space.

        # We can check if the constrained value differs from the unconstrained value.

        bounded_params = set()

        if constrained_samples:

            for param_name, z_array in z_values.items():

                if param_name in constrained_samples:

                    c_arr = constrained_samples[param_name]

                    # If the values are significantly different, it's likely a bounded prior

                    # (e.g., z=100 maps to c=1.0)

                    if hasattr(c_arr, 'shape') and c_arr.size > 0:

                        # Compare first walker's z and c values

                        z_val = float(z_array[0])

                        c_val = float(c_arr[0, 0] if c_arr.ndim >= 2 else c_arr[0])

                        if abs(z_val - c_val) > 1e-3:

                            bounded_params.add(param_name)

        # Collect all z values first to track stats

        for param_name, z_array in z_values.items():

            for i in range(num_walkers):

                z_val = float(z_array[i])

                if jnp.isnan(z_val):

                    walker_stats[i]['has_nan'] = True

                elif jnp.isinf(z_val):

                    walker_stats[i]['has_inf'] = True

                elif param_name in bounded_params:

                    if walker_stats[i]['max_z_bounded'] is None:

                        walker_stats[i]['max_z_bounded'] = abs(z_val)

                    else:

                        walker_stats[i]['max_z_bounded'] = max(walker_stats[i]['max_z_bounded'], abs(z_val))

        # Format header and print parameter table with constrained/unconstrained side-by-side

        col_width = 35

        header = f"  {'Parameter':<20s}"

        for i in range(num_walkers):

            header += f" | {'W'+str(i)+' (Uncons / Constr)':<{col_width}}"

        separator = f"  {'-'*20}" + "".join([f" | {'-'*(col_width)}" for _ in range(num_walkers)])

        log_func(plModel, header)

        log_func(plModel, separator)

        # Rows for each sampling parameter

        for param_name, z_array in z_values.items():

            row_vals = []

            for i in range(num_walkers):

                z_val = float(z_array[i])

                # Get constrained value if available

                c_val_str = "N/A"

                if constrained_samples and param_name in constrained_samples:

                    c_arr = constrained_samples[param_name]

                    if hasattr(c_arr, 'shape') and c_arr.ndim >= 2:

                        c_val = float(c_arr[i, 0])

                        c_val_str = f"{c_val:8.6f}"

                    elif hasattr(c_arr, 'shape') and c_arr.ndim == 1:

                        if i < len(c_arr):

                            c_val = float(c_arr[i])

                            c_val_str = f"{c_val:8.6f}"

                row_vals.append(f"{z_val:8.6f} / {c_val_str}")

            row_str = " | ".join([f"{val:<{col_width}}" for val in row_vals])

            log_func(plModel, f"  {param_name:<20s} | {row_str}")

        # Compute potential energy on-the-fly for each walker

        log_func(plModel, separator)

        potential_energy_array = None

        grad_norms = [0.0] * num_walkers

        if prob_fn is None:

            # Try to extract probability function from kernel

            try:

                prob_fn = mcmc_obj.kernel.model_fn

            except AttributeError:

                log_func(plModel, f"  {'potential_energy':<20s} | {'N/A':<{col_width}}" + " | " + " | ".join([f"{'N/A':<{col_width}}" for _ in range(num_walkers-1)]))

                log_func(plModel, f"  {'gradient_norm':<20s} | {'N/A':<{col_width}}" + " | " + " | ".join([f"{'N/A':<{col_width}}" for _ in range(num_walkers-1)]))

                prob_fn = None

        if prob_fn is not None:

            try:

                from numpyro.infer.util import potential_energy as _numpyro_pe

                # potential_energy accepts unconstrained params directly and returns

                # -log_joint (including change-of-variables Jacobian terms), which

                # is exactly the quantity HMC/ensemble samplers minimise.

                def log_prob_fn(z_dict):

                    return -_numpyro_pe(prob_fn, (plData, plModel), {}, z_dict)

                potential_energy_array = []

                grad_arrays = {param_name: [] for param_name in z_values.keys()}

                for i in range(num_walkers):

                    walker_z = {k: v[i] for k, v in z_values.items()}

                    log_prob, grads = jax.value_and_grad(log_prob_fn)(walker_z)

                    potential_energy_array.append(-float(log_prob))

                    for param_name in z_values.keys():

                        grad_val = grads.get(param_name, 0.0)

                        grad_arrays[param_name].append(float(grad_val))

                        grad_norms[i] += float(grad_val) ** 2

                grad_norms = [jnp.sqrt(g) for g in grad_norms]

            except Exception as e:

                log_func(plModel, f"  Warning: Could not compute potential energy on-the-fly: {e}")

                potential_energy_array = None

        # Log potential energy

        if potential_energy_array is not None:

            row_vals = [f"{pe_val:8.4e}" for pe_val in potential_energy_array]

            row_str = " | ".join([f"{val:<{col_width}}" for val in row_vals])

            log_func(plModel, f"  {'potential_energy':<20s} | {row_str}")

        else:

            row_vals = ["N/A"] * num_walkers

            row_str = " | ".join([f"{val:<{col_width}}" for val in row_vals])

            log_func(plModel, f"  {'potential_energy':<20s} | {row_str}")

        # Log gradient norms

        log_func(plModel, separator)

        if grad_norms and any(g > 0 for g in grad_norms):

            row_vals = [f"{gn:8.4e}" for gn in grad_norms]

            row_str = " | ".join([f"{val:<{col_width}}" for val in row_vals])

            log_func(plModel, f"  {'gradient_norm':<20s} | {row_str}")

        else:

            row_vals = ["N/A"] * num_walkers

            row_str = " | ".join([f"{val:<{col_width}}" for val in row_vals])

            log_func(plModel, f"  {'gradient_norm':<20s} | {row_str}")

        # === WARNINGS ===

        log_func(plModel, "\n" + "="*150)

        log_func(plModel, "=== INITIALIZATION DIAGNOSTICS ===")

        log_func(plModel, "="*150)

        warnings = []

        # Check for NaN/Inf in z values

        for i in range(num_walkers):

            if walker_stats[i]['has_nan']:

                warnings.append(f"⚠️  Walker {i}: Contains NaN in unconstrained space (z). Model may be evaluating outside support.")

            if walker_stats[i]['has_inf']:

                warnings.append(f"⚠️  Walker {i}: Contains Inf in unconstrained space (z). Walkers may have jumped to boundaries.")

        # Check for extreme z values (indicates sigmoid saturation)

        for i in range(num_walkers):

            if walker_stats[i]['max_z_bounded'] is not None and walker_stats[i]['max_z_bounded'] > 100:

                warnings.append(f"⚠️  Walker {i}: Extremely large |z| values ({walker_stats[i]['max_z_bounded']:.2e}) for bounded parameters. "

                               f"Sigmoid will saturate to exact boundaries (0.0 or 1.0). "

                               f"Indicates poor prior-likelihood alignment or aggressive ensemble moves.")

        # Check for massive potential energy

        if potential_energy_array is not None:

            for i in range(num_walkers):

                pe_val = potential_energy_array[i]

                if pe_val > 1e10:

                    warnings.append(f"🔴 Walker {i}: CRITICAL potential energy = {pe_val:.4e}. "

                                   f"Prior and likelihood are severely misaligned. "

                                   f"Check priors encompass posterior region.")

                elif pe_val > 1e6:

                    warnings.append(f"⚠️  Walker {i}: Very high potential energy = {pe_val:.4e}. "

                                   f"Walker started in low-probability region.")

        # Check for extremely large gradients

        if grad_norms:

            for i in range(num_walkers):

                if len(grad_norms) > i and grad_norms[i] > 1e8:

                    warnings.append(f"⚠️  Walker {i}: Extreme gradient norm = {grad_norms[i]:.4e}. "

                                   f"Indicates numerical instability or posterior discontinuity.")

        # Check if all walkers are at identical positions (suggests init_to_median without diversification)

        all_same = True

        first_z = {k: float(v[0]) for k, v in z_values.items()}

        for i in range(1, num_walkers):

            for k, v in z_values.items():

                if abs(float(v[i]) - first_z[k]) > 1e-6:

                    all_same = False

                    break

            if not all_same:

                break

        if all_same:

            warnings.append(f"⚠️  ALL WALKERS INITIALIZED AT IDENTICAL POSITION. "

                           f"Using init_to_median without randomness? Ensemble will spend first N steps diversifying.")

        # Print warnings or success message

        if warnings:

            for warning in warnings:

                log_func(plModel, warning)

        else:

            log_func(plModel, "✓ Initialization looks reasonable. No critical issues detected.")

        log_func(plModel, "="*150 + "\n")

    else:

        log_func(plModel, f"  z_values is not a dict: {z_values}")

def _log_warmup_energy_summary(final_energies, num_chains, plModel, logger):

    """Quick energy table printed immediately after warmup"""

        # The same kernel instance is reused for Phase 3b ensemble production

        # (if multimodal fallback is triggered in Phase 2).

        ensemble_kernel = make_ensemble_kernel(warmup_sampler)

        # init_to_uniform(radius=2) draws each walker uniformly in unconstrained

        # space within [-2, +2], which maps to the interior of every prior's

        # support (for Uniform(lo,hi) this covers roughly the central 76%).

                             num_chains=num_chains,

                             progress_bar=False, chain_method='vectorized')

        # Reverting the failing line.

        if existing_warmup_state is not None and existing_warmup_done >= num_warmup:

            # ── SKIP: warmup already complete in checkpoint ──────────────────

            # Walker positions saved by the previous run's warmup_complete checkpoint

            ensemble_mcmc.run(jax.random.PRNGKey(random_seed + existing_warmup_done + 1),

                              pl_data=plData, pl_model=plModel)

            jax.block_until_ready(ensemble_mcmc.last_state)

            _warmup_jit_elapsed = ttt.time() - _warmup_pilot_start

            _log(plModel, f"  JIT compilation: {_warmup_jit_elapsed:.1f}s")

                _log(plModel, f"=== PHASE 1: {warmup_sampler.upper()} ENSEMBLE WARMUP ===")

                _log(plModel, f"Walkers: {num_chains}, Warmup iterations: {num_warmup}")

                # from .sampler_utils import check_walker_initialization as _cwi

                # _cwi(model_fn,plData,plModel,random_seed,_log)

            # ── JIT-trigger pilot step ────────────────────────────────────────

            # Run exactly 1 sample to force XLA compilation before the main

            # loop.  This step is needed for both fresh starts and resumes because

            ensemble_mcmc.run(jax.random.PRNGKey(random_seed + _warmup_start_iter),

                              pl_data=plData, pl_model=plModel)

            jax.block_until_ready(ensemble_mcmc.last_state)

            if existing_warmup_state is  None or existing_warmup_done == 0:

                # Log initial walker positions using the helper

                from .logging import check_walker_initialization as _cwi

                _cwi(ensemble_mcmc, plModel, _log, prob_fn=model_fn, plData=plData)

            _warmup_jit_elapsed = ttt.time() - _warmup_pilot_start

            _log(plModel, f"  JIT compilation + 1 pilot step: {_warmup_jit_elapsed:.1f}s  "

                 f"— warmup rate below reflects true hardware throughput")

import time as ttt

import numpy as np

import jax

import jax.numpy as jnp

import numpyro

from .checkpoint import save_pickle

def run_batched(mcmc_obj, target, chain_count, plData, plModel, checkpoint_interval,

                save_path, random_seed, num_chains, requested_num_chains, num_warmup,

                device_type, sampler_type, param_config, config, format_time,

                _format_rate, log_func, start_offset=0, existing_grouped=None,

            log_func(plModel, f"ESS moves (from config): {{{', '.join(log_entries)}}}")

        else:

            moves = {ESS.DifferentialMove(): 0.7, ESS.KDEMove(): 0.3}

        return ESS(model_fn, moves=moves, init_strategy=init_to_uniform(radius=ensemble_init_radius))

        # return ESS(model_fn, moves=moves, init_strategy=init_to_uniform(radius=ensemble_init_radius))

        from numpyro.infer import init_to_median

        return ESS(model_fn, moves=moves, init_strategy=init_to_median())

      <input type="number" id="m-nw" value="1000" min="1">

      <div class="help">Number of ensemble warmup iterations.</div>

    </div>

    <div class="fg">

      <label><span class="key-name">ensemble_init_radius</span> <span class="badge b-opt">Optional</span></label>

      <input type="number" id="m-eir" value="2.0" step="0.1">

      <div class="help">Initialisation radius for ensemble walkers in unconstrained space. Default: <code>2.0</code>. Set to &lt; 2 (e.g. <code>0.5</code>) to avoid walkers being initialized exactly at prior boundaries (<code>lo</code>), which can stall walkers in bounded priors.</div>

    </div>

  </div>

  <!-- ── walkers: auto, ess, aies ── -->