Fix GPU-only NaN gradient in form factor pole integral

[joglekara]

Problem

Running `run_tsadar.py --mode fit` on a GPU produces a NaN gradient in `grad.electron.normed_Te` after the first epoch. The NaN feeds back into the optimizer (adam), making every subsequent epoch's loss NaN, and the run eventually crashes. The same fit runs fine on CPU.

Reproduced off `main` on a Perlmutter A100 (bundled shot 111411; the defect is shot-independent).

Root cause

A `jnp.where` NaN-gradient leak in the form factor's principal-value integral, `ratcen()` in tsadar/core/physics/ratintn.py:

```python
rf = fav / gav + tmp * gdif / (12.0 * gav3) # selected far from a pole
rfn = fdif / gdif + tmp * jnp.log(...) / gdif2 # selected near a pole
out = jnp.where(|gdif| < 1e-4*|gav|, rf, rfn)
```

Near a pole, `gav = ½(g[i] + g[i-1]) → 0`, so the unselected `rf` branch is `inf`. `jnp.where` returns the correct value (`rfn`), but reverse-mode autodiff differentiates both branches and propagates `0 × (inf gradient of rf)`:

• CPU: `gav` lands at a tiny non-zero (~1e-16) → `1/gav³` is huge but finite → `0 × huge = 0`. Harmless.
• GPU: FMA/rounding lands `gav` at exactly `0` → `1/gav³ = inf` → `0 × inf = NaN`.

The NaN flows back through `vTe = sqrt(Te/Me)` → `grad.electron.normed_Te`.

Fix

Apply the double-where safe-divide idiom: guard `gav` in the `rf` branch (`gav_safe = jnp.where(use_rf, gav, 1.0)`) so the dead branch's gradient stays finite. The selected value is unchanged — the forward pass is numerically identical and the fix is backend-independent.

Also pre-initialize `best_weights` in `1d_optax_loop` so a never-improving (e.g. NaN) batch returns valid params instead of raising `UnboundLocalError` (a secondary crash that masked the real bug). Defensive; can be dropped if a minimal diff is preferred.

Verification

Full GPU fit run on Perlmutter A100 completes with `EXIT=0`, zero NaN/errors, finite decreasing losses across all batches, through postprocessing. Previously NaN at epoch 2.

🤖 Generated with Claude Code