Changes for .kilo/runner_refactor_plan.md: 82 added lines, 0 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.
Changes for docs/rajpaul_bis_analysis.md: 7 added lines, 2 removed lines.
Original line number
Diff line number
Diff line
@@ -400,10 +400,15 @@ Given that you are working with CHEOPS data and your own RV+ancillary-indicator
**Step 1: Quick correlation-based screening (before full model fit)**
1. Compute the measured Pearson correlation $\rho$ between your RV and BIS time series.
2. From a preliminary GP fit or spectral analysis, estimate the realized variances $\sigma_g^2$ and $\sigma_{\dot g}^2$ in your baseline. Check that both activity timescales (rotation modulation and spot evolution) are well-represented — if one is barely visible, flag this as data-driven rank deficiency.
2. Estimate the realized variances $\sigma_g^2$ (rotation modulation) and $\sigma_{\dot g}^2$ (spot evolution) in your baseline. Since these represent the power in the two basis functions, you can obtain them via:
-**Preliminary Single-Indicator GP Fit (Recommended):** Run a quick GP fit on RV alone using a Quasi-Periodic kernel. While this fits the observable rather than the latent process, you can use the resulting GP variance as a proxy for $\sigma_g^2$. (Note: The inferred GP amplitude is actually the product of the coupling coefficient and the latent variance, i.e., $\text{Var}(\text{RV}_{\text{GP}}) \approx V_c^2 \sigma_g^2$. However, because the determinant $\Delta$ also scales proportionally with these coupling coefficients, the absolute scales cancel out when evaluating the system's conditioning relative to the total observable variance. This makes the screening result scale-invariant and valid for a first-order estimate). To estimate $\sigma_{\dot g}^2$, use the "unit projection" approximation: $\sigma_{\dot g}^2 \approx \sigma_g^2 / \ell^2$ (where $\ell$ is the decay lengthscale). This assumes that the observable is dominated by the rotation modulation term and that the derivative term's power scales inversely with the square of the evolution timescale, providing a sufficient order-of-magnitude estimate for screening.
-**Spectral Analysis:** Use a Lomb-Scargle periodogram. Integrate power under the main rotation peak for $\sigma_g^2$ and the broadband/wing power for $\sigma_{\dot g}^2$.
-**Stellar Priors:** Use known $v \sin i$, $R_*$, and typical spot lifetimes for the stellar type as a first-order approximation.
*Check that both activity timescales are well-represented — if one is barely visible (e.g., $\ell$ hits the prior boundary in a GP fit), flag this as data-driven rank deficiency.*
4. If $|\Delta|$ is small, proceed cautiously — you have evidence of near-degeneracy. If $|\Delta|$ is large, RV and BIS are accessing different directions; proceed with confidence.
4. If $|\Delta|$ is small compared to the scale set by the product $\sigma_g\sigma_{\dot g}$ or its posterior uncertainty, proceed cautiously — you have evidence of near-degeneracy. If $|\Delta|$ is large on that scale, RV and BIS are accessing different directions; proceed with confidence.
**Step 2: Full posterior-based validation (after joint model fit)**
