Changes for .kilo/runner_refactor_plan.md: 0 added lines, 82 removed lines.
Original line number
Diff line number
Diff line
# PyAETNA `runner.py` Refactor Plan
## Goal
Create a focused refactor of `PyAETNA/sampling/engine/runner.py` to:
1. remove redundant or duplicate code,
2. remove dead variables,
3. remove useless comments,
4. fix incorrect logic,
5. improve performance.
---
## 1. Remove Redundant / Duplicate Code
- Eliminate duplicated config reads and overriding assignments:
- Remove the second `checkpoint_interval = mcmc_config.get('checkpoint_interval', 100)`.
- Remove the second `max_tree_depth = mcmc_config.get('max_tree_depth', 6)`.
- Remove the second `target_accept_prob = mcmc_config.get('target_accept_prob', 0.9)`.
- Keep the one validated assignment for each value and drop the later fallback read.
- Replace the early `num_chains_if_nuts = mcmc_config_peek.get('num_chains_if_nuts', 4)` with a separate hint variable, e.g. `num_chains_if_nuts_hint`, and use it only for early device-count estimation.
- Remove the unnecessary intermediate `ensemble_mcmc` object in direct ESS/AIES production mode. Use `_unified_ensemble_state` directly for `batch_mcmc.post_warmup_state` and `batch_mcmc._warmup_state`.
- Move `dense_mass = mcmc_config.get('dense_mass', False)` into the validated config-read block for consistency.
## 2. Remove Dead Variables
- Replace all unused `_jit_time` unpackings from `run_batched(...)` with `_`.
- Replace `warmup_energy = np.zeros((num_chains, num_warmup))` with a smaller representation, because only first and last warmup energies are used.
- Eliminate the redundant `samples_grouped = {}` reset when a checkpoint is incompatible.
- Inline `num_cores` into its log statement, since it is only used once.
- Remove `ensemble_mcmc.last_state` assignment in direct ensemble mode if it is not guaranteed to be populated.
- Replace `_pilot_elapsed = _unified_jit_elapsed if _unified_jit_elapsed else 0.0` with a direct expression where used.
## 3. Remove Useless Comments
- Remove overly verbose or redundant comments such as:
- configuration file path resolution notes that merely restate the code
-`reset was already determined above`
-`Just return the existing data`
-`Reverting the failing line.`
- repeated explanatory blocks about constrained-space samples and postprocess behavior
-`Note: Raw sample summary is now printed in reports.py` unless it is needed for developer context
- Preserve only comments that document non-obvious behavior and control-flow decisions.
## 4. Fix Incorrect Logic
- Fix the `NameError` bug in the dense mass logging block: replace `axis1= unfolded` with `axis1=1`.
- Fix the warmup SKIP branch so it runs a single pilot `ensemble_mcmc.run(...)` before calling `ensemble_mcmc.get_samples(group_by_chain=True)`.
- Fix direct ensemble production so `batch_mcmc.post_warmup_state` uses `_unified_ensemble_state` directly instead of `ensemble_mcmc.last_state`.
- Ensure `checkpoint_interval`, `max_tree_depth`, and `target_accept_prob` are not silently overwritten after validation.
- Adjust warmup pilot RNG handling so it does not reuse the same key on resumed warmup runs.
## 5. Improve Performance
- Replace repeated warmup `np.concatenate` growth with batch accumulation lists and a single final `np.concatenate`.
- Consolidate repeated NUTS adaptation logging into a helper function to remove duplication and keep logging bug fixes in one place.
- Consider batching the Phase 2 energy computation with `jax.vmap` if `_compute_log_joint` can be vectorized, avoiding a Python loop over chains.
- Remove unnecessary repeated late imports or move them to the module-level import block where safe.
- Document the JAX device-count over-provision strategy so future maintainers understand the pre-initialization constraint.
---
## Implementation Notes
- Prefer smaller, atomic changes in the refactor:
1. fix config duplicate-read logic,
2. repair warmup SKIP/run state handling,
3. remove dead variables and comments,
4. consolidate repeated NUTS logging,
5. optimize warmup sample accumulation.
- Keep the existing structure of the auto/NUTS/ensemble pipeline intact while removing the redundant blocks that cause bugs and confusion.
- Validate the file with `python3 -m py_compile PyAETNA/sampling/engine/runner.py` after each major change.
---
## Success Criteria
-`runner.py` contains no duplicate validated config reads.
- No unused local variables remain in the hot code path.
- The warmup SKIP/resume paths no longer fail on `get_samples()` with `NoneType` state.
- Dense mass logging no longer throws `NameError`.
- The direct ensemble production path preserves the JIT-initialized state correctly.
- Warmup accumulation is more efficient and avoids O(N²) concatenation overhead.
