Changes for PyAETNA/.kilo/boundary_init_analysis.md: 58 added lines, 0 removed lines.
Original line number
Diff line number
Diff line
# Analysis: Boundary Initialization Bug in NumPyro `init_to_uniform`
## Problem Description
In PyAETNA, when using the `auto` sampling pipeline, the ensemble warmup phase (ESS/AIES) is configured to use `init_to_uniform(radius=2)`. This strategy is intended to draw initial samples in the unconstrained space (typically $[-2, 2]$) and then transform them back to the constrained space of the priors. This ensures that the starting points are strictly inside the prior supports, avoiding numerical instabilities (like $\text{logit}(0) = -\infty$) that occur at the boundaries.
However, evidence from `pyaetna_config.log` indicates that for **bounded priors** (including `TruncatedNormal`), the initial values (`_first`) are being set **exactly** at the lower or upper boundaries of the prior. This suggests the problem is not limited to a specific distribution but is systemic to any prior with hard constraints.
### Evidence from Logs
In `pyaetna_config.log`, we observe:
-`prior__stellar_density` with bounds $[0.07, 0.14]$ initializes at exactly `0.07` or `0.14`.
-`prior__b`, `prior__period`, and `norm` similarly show initialization at their respective prior limits.
When a value is exactly on the boundary of a constrained distribution, its representation in unconstrained space becomes $\pm\infty$. This leads to:
1.**Numerical Instability**: NaNs during the first few iterations.
2.**Slow Warmup**: The sampler struggles to move away from the boundary, significantly increasing the time required for the ensemble phase to converge or causing it to fail and trigger the NUTS fallback.
## Technical Context
### PyAETNA Implementation
The kernel factory in `PyAETNA/sampling/engine/sampler_utils.py` implements the initialization as follows:
```python
# Conceptual snippet of the current implementation
The bug likely resides within NumPyro's `init_to_uniform` implementation or its interaction with `TruncatedNormal`. Instead of mapping the $[-2, 2]$ range to the interior of the support, it may be clipping values or failing to correctly transform the boundaries for truncated distributions, resulting in values that land exactly on the edge.
## Investigation Plan
### 1. Isolated Reproduction
Create a minimal standalone script that:
- Defines a model with a `dist.TruncatedNormal` prior.
- Uses `MCMC` with `init_strategy=init_to_uniform()`.
- Extracts the initial values of the parameters before the first sampling step.
- Verifies if these values match the `low` or `high` bounds of the distribution.
*Note: Previous attempts failed due to `EnsembleSampler` requiring `num_chains > 1` and `chain_method='vectorized'`. The reproduction script must be updated to meet these requirements.*
### 2. Version Correlation
- Compare the behavior across different NumPyro versions (e.g., current vs. latest 0.21.0).
- Check NumPyro's GitHub issues for "init_to_uniform" or "TruncatedNormal boundary" reports.
### 3. Alternative Initialization Testing
Test alternative strategies to see if they avoid the boundary issue:
-`init_to_median()`: Does it start safely inside?
- Custom initialization: Draw from a "shrunk" version of the prior (e.g., $low + 0.1 \times (high-low)$).
### 4. Proposed Fix
If the bug is confirmed in NumPyro, implement a custom initialization wrapper in PyAETNA that ensures samples are drawn from the interior of the support, bypassing the faulty `init_to_uniform` for constrained parameters.
Changes for PyAETNA/.kilo/eliminate_safeguards_plan.md: 31 added lines, 36 removed lines.
Original line number
Diff line number
Diff line
# Strategy: Eliminating Model Safeguards and Fixing Boundary Initialization
# Strategy: Eliminating Model Safeguards for Numerical Honesty
## 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**.
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 systemic numerical instabilities and create "invisible walls" in the parameter space.
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.
### The Criticality of Safeguards:
1.**Gradient Masking**: Safeguards like clipping create flat plateaus in the likelihood surface. When a sampler hits a clipped region, the gradient becomes zero, leaving the walker stranded with no signal to guide it back toward the high-probability region.
2.**Hidden Failures**: By suppressing `NaN`s, the model hides legitimate failures (e.g., parameters moving into physically impossible regimes). This makes debugging difficult, as the sampler may appear to be converging or "stuck" without any explicit error signal.
3.**Numerical Instability**: Masking errors in the forward pass can lead to unstable updates in the sampler, resulting in 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)**.
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 and 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.
### Step 1: Audit and Mapping
Identify every instance of numerical safeguarding in the physics engine.
*This ensures that we cannot start a run that is mathematically doomed to fail due to $z = \pm\infty$.*
**Audit Findings (`body_models.py`):**
### Step 2: Hardening the Priors
To make the model safe without internal clipping, we shift the boundaries of the priors themselves.
| Function | Safeguard | Physical Reason / Criticality |
| :--- | :--- | :--- |
| `compute_inclination` | `jnp.clip(..., 0.0, 1.0 - 1e-10)` | Prevents `arccos` from receiving values outside $[-1, 1]$. Critical if $b$ or $a/R_*$ are poorly constrained. |
| `compute_transit_duration` | `jnp.where(transit_occurs, b_param, 0.0)` | `b_safe` prevents `sqrt` of negative numbers when no transit occurs. Masks the fact that the walker is in a non-transiting regime. |
| `compute_transit_duration` | `jnp.maximum(k_eff**2 - b_safe**2, 0.0)` | Flooring for `sqrt`. Redundant if `b_safe` is used, but creates a flat gradient at the transit boundary. |
| `compute_transit_duration` | `jnp.where(transit_occurs, duration, 0.0)` | Forces duration to 0 if no transit. Creates a discontinuous jump in the likelihood surface. |
| `compute_eclipse_duration` | `jnp.where(eclipse_occurs, b_param, 0.0)` | `b_safe` for secondary eclipse. Same as transit: masks non-eclipse regimes. |
| `compute_eclipse_duration` | `jnp.maximum(k_eff**2 - b_safe**2, 0.0)` | Flooring for `sqrt`. Same as transit. |
| `compute_eclipse_duration` | `jnp.where(eclipse_occurs, duration, 0.0)` | Forces duration to 0. Same as transit. |
| `compute_phase_angle_fast` | `jnp.clip(-z / r_norm, -1.0, 1.0)` | Prevents `arccos` domain errors. Critical if $z$ or $r_{norm}$ are unstable. |
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`.
**Summary:**
Most safeguards are protecting `arccos` and `sqrt` from domain errors. The `jnp.where` logic in duration functions is particularly problematic as it creates "dead zones" with zero gradients for any walker that drifts outside the transit/eclipse condition.
### Step 3: Fixing the Initialization Root Cause
Address the bug in the interaction between `init_to_uniform` and bounded distributions.
### Step 2: Hardening the Input Stage
Ensure that the inputs reaching the physics engine are guaranteed to be within the valid domain, removing the need for internal clipping.
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.
**Proposed Solutions:**
1.**$\epsilon$-Shrinking of Priors**: Automatically shrink bounded priors by a small margin $\epsilon$ (e.g., $10^{-6}$) to ensure values never hit the exact mathematical boundary.
2.**Tighter Meaningful Priors**: Use physically motivated, tighter priors to prevent walkers from exploring regions where the model is numerically unstable.
3.**Input Validation**: Implement a validation check before the sampling loop to ensure initialized walkers are in a safe regime.
### Step 4: Purging the Safeguards
### Step 3: Purging the Safeguards
Once the input stage is guaranteed to be safe, remove the "crutches" from the physics engine.
**Action Items:**
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.
-**Faster Convergence**: By removing flat plateaus, the sampler will have healthy gradients across the entire valid parameter space.
-**Cleaner Code**: The physics models in `body_models.py` will be stripped of boilerplate safety logic, making them easier to maintain and read.
